Converting a tensorflow model to jax

[PabloAMC]

Hi!
Is there a way to convert a loaded tensorflow model to jax? I have replicated a paper architecture, originally in Tensorflow, in Jax, but would love to be able to load the .pb file for the weights. Is there a way to do so?
Thanks a lot in advance

[8bitmp3]

Hi @PabloAMC Have you come across this library https://github.com/deepmind/tf2jax ?

[8bitmp3]

From the README's Support for Serialization Formats section:

model = tf.Module()
model.f = forward
model.f(x)  # Dummy call.
tf.saved_model.save(model, "/tmp/blah")

restored = tf.saved_model.load("/tmp/blah")
jax_func, jax_params = tf2jax.convert(restored.f, np.zeros_like(x))

[8bitmp3]

Also mentioned in the tf2jax README:

jax2tf now also offers the experimental call_tf function which allows JAX to call TensorFlow functions. For compiled code, this works by staging out TensorFlow functions to XLA.

And from the jax2tf README: JAX and TensorFlow interoperation (jax2tf/call_tf) (cc author @gnecula)

jax2tf.call_tf: for using TensorFlow functions in a JAX context, e.g., to call a TensorFlow library or a SavedModel inside a JAX function.
...
Since the TensorFlow functions that are produced by jax2tf.convert can be JIT-compiled by design, we can call them using jax2tf.call_tf thus achieving a round-trip from JAX to TensorFlow (e.g., a SavedModel) and back.

Further down the doc, there's a section called Calling TensorFlow functions from JAX with more details (and limitations).

[gnecula]

If you have already replicated the model in JAX, what kind of conversion of the TF model to JAX are you looking for? tf2jax and jax2tf.call_tf would not help with loading the weights; they would allow you to run the TF model in a JAX context.