Query: Availability of JAX-based Implementations for e3nn.nn.models.gate_points_2101.Compose Module #66
-
|
Hello everyone, I hope you're doing well. I'm currently exploring the capabilities of the E3NN library and am particularly interested in leveraging JAX for improved performance and compatibility with accelerators like GPUs and TPUs. I'm curious to know if there are any existing implementations or ongoing efforts to implement the Compose module from the e3nn.nn.models.gate_points_2101 submodule using JAX. I believe such implementations could provide significant benefits in terms of performance and scalability. If anyone has information about this or has worked on similar endeavors, I would greatly appreciate your insights and guidance. Additionally, if there are any relevant resources or documentation available, please feel free to share them. Thank you in advance for your help and contributions. |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment
-
|
You have this module https://github.com/e3nn/e3nn-jax/blob/main/e3nn_jax/experimental/point_convolution.py That you can compose with an e3nn_jax.gate operation. But no, I don't have a drop in replacement for e3nn.nn.models.gate_points_2101. You have to reimplement it |
Beta Was this translation helpful? Give feedback.
You have this module
https://github.com/e3nn/e3nn-jax/blob/main/e3nn_jax/experimental/point_convolution.py
That you can compose with an e3nn_jax.gate operation.
But no, I don't have a drop in replacement for e3nn.nn.models.gate_points_2101. You have to reimplement it