Porting models based on e3nn-jax for deployment

[rodhakate]

Is there a way and is it possible to "correctly (without breaking any equivariant parts)" convert models based on e3nn-jax to TFLite?

Convert JAX --> TensorFlow via jax2tf,
Convert SavedModel --> TFLite/LiteRT.

Thank you in advance.

[teddykoker]

Hi @rodhakate have you looked into this guide? Happy to look into this but I am curious if you have tried this and run into any issues.

[rodhakate]

Thank you for the reply.
Actually I was using the torch version of e3nn library, but later realized it wont be straightforward to port a model (using torch e3nn) to onnx or tf.

So I started implementing c++ library for e3nn, actually building on your e3nn.c repo since it dosent have all the functionalities of the official e3nn such as FulyConnectedTensorProduct (FCTP) and others.

It was just getting too complex so I decided to replace FulyConnectedTensorProduct with tensor product (TP) + linear, which is easier to implement in c++.

Also they should yield similar results since,

for FCTP the layer owns the weights and are learned during training and then are fixed at inference.
TP + Linear has no weights inside. Another small MLP produces the weights each forward pass, so they can change per edge/time/sample.

So same kind of equivariant mixing, but the only difference would be who owns the weights (inside the layer vs. an external MLP) and whether they’re static or input-conditioned.

I am yet to validate the c++ implementation of TP + Linear.

And recently I realized if I implement my model using e3nn-jax then maybe it is possible to port it to TF in a straightforward way, but wasn't sure. I will look into "https://ai.google.dev/edge/litert/models/jax_to_tflite".

Thanks for the info.

Also If i get the c++ imeplementation of TP + Linear working I will post it here.

[mitkotak]

Hey have you looked at torch.export + AOTI ? We use it in Nequip/Allegro to package and compile into a .so which is then loaded up in LAMMPS