-
Notifications
You must be signed in to change notification settings - Fork 45
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Jax support #55
Comments
@Gabri95 ? |
I am sorry for the late reply, but I was on vacation in the last couple of weeks. I am not currently planning to support other frameworks like Jax. I am not very familiar with Jax, so it is hard for me to precisely estimate how much work it would be, but I'd be very happy to provide some support if an user would like to try doing that! Best, |
Thanks @Gabri95! |
Hey @Gabri95, |
@Gabri95 |
Thanks a lot for porting the library in Jax! The I need to push some updates on the library in the coming week (sorry for being a little offline recently :/ ) and I will also include a pointer to your code in the documentation. Thanks a lot again! Gabriele |
Happy to eventually works ahah was more than expected and there are still quite some layers not supported, but should be easy to add them taking inspiration from the ones I've ported :)
|
This is a good question, let me ask for some advice about this and come back when I know more, sorry Regarding pypi, you can follow these intructions (but I'd recommend trying with testpypi first) Gabriele |
I've just included a reference to your fork in the README of the library! Regarding where you should host the repo, it seems possible to both keep it where it is now and move it under QUVA-Lab. Again, I don't have a strong preference, but maybe we can keep it under your account for the moment and move it to quva-lab in the future. Thanks again for this amazing work! 😄 Best, |
Hi,
Thanks for this really cool library!
Being a jax user, I was wondering whether you've thought on extending it to support jax, akin to e3nn-jax?
At first glance it seems that most of the library is 'pure' python, appart from the
GeometricTensor
andFieldType
classes inescnn/nn/*
, which seems easily translatable to jax.Array etc, and obviously all the layers inescnn/nn/modules
which would need to be rewritten for flax / haiku / equinox.I'd be happy to help out with this :)
Best,
Emile
The text was updated successfully, but these errors were encountered: