How to export the nnx Module to SavedModel?

[5c4lar]

I'm trying to load and use the model trained with flax nnx with C++. jax2tf seems to provide a way to achieve this, but the examples are about linen. How can we convert the flax nnx module to tf?

[cgarciae]

Here an example using nnx.split and nnx.merge to use jax2tf to export an Flax NNX Module as a SavedModel:

from flax import nnx
import jax
import jax.numpy as jnp
from jax.experimental import jax2tf
import tensorflow as tf

model = nnx.Linear(3, 4, rngs=nnx.Rngs(0))

graphdef, state = nnx.split(model)
state = jax.tree.leaves(state) # flatten the state

tf_state = tf.nest.map_structure(tf.Variable, state)

def forward_jax(state, x):
  model = nnx.merge(graphdef, state)
  return model(x)

# test forward
y = forward_jax(state, jnp.ones((3,)))

def predict_tf(x):
  return jax2tf.convert(forward_jax)(tf_state, x)

tf_model = tf.Module()
# Tell the model saver what are the variables.
tf_model._variables = tf.nest.flatten(tf_state)
tf_model.f = tf.function(predict_tf, jit_compile=True, autograph=False)
tf.saved_model.save(tf_model,'some/path')

Edit: updated to use more recent APIs

[5c4lar]

Thanks! I didn't notice the to_pure_dict method when I tried to implement this myself, this example is very helpful.