Autodifferentiating through a model exported from Flax then requires the conversion of a traced array

[Peter-Vincent]

I have a model I am trying to optimise. A component of that model is a neural network. That neural network is trained beforehand (since it can be used by a variety of other models), saved using the following code

jax_module = JaxModule({"params":state.params},state.apply_fn,trainable=True,input_polymorphic_shape='(b, ... )')
        em = ExportManager(jax_module, [
        ServingConfig(
            'serving_default',
            input_signature=[
                tf.TensorSpec([num_rands,1, 2], tf.float32, name='inputs')
            ],
            tf_postprocessor=lambda x: dict(outputs=x)),
        ])
        model_path = os.path.join(PRE_TRAINED_DIR,f"model_save_path")
        em.save(model_path)

And then loaded in when I need it with the following code

loaded_model = tf.saved_model.load(model_path)

This is following the documentation and examples in the Orbax repository

The problem is arising now when I try and optimise this second model, for which the neural network is a small component.

If I just run the model without wrapping it in grad (or value_and_grad) it works fine. However, if instead the function is wrapped by value_and_grad, then I get these errors

TypeError: Binding inputs to tf.function failed due to `The numpy.ndarray conversion method __array__() was called on traced array with shape float32[36878,1,2].

when I try and call the model with

model(inputs)["outputs"]

I understand at a high level why this is happening, but I don't see a way around this. This will also be a problem if this line were to succeed, since the outputs of the model are tf.float32 arrays, which jax.numpy functions like (for example) jnp.mod can't process, so the arrays need to be cast back into recognisable types.

Any recommendations on how to resolve this would be greatly appreciated!

Many thanks