jax2tf Examples
===============

Link: go/jax2tf-examples.

This directory contains a number of examples of using the
[jax2tf converter](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md) to:

  * save SavedModel from trained MNIST models, using both Flax and pure JAX.
  * reuse the feature-extractor part of the trained MNIST model
    in a larger TensorFlow Keras model.
  * use Flax models with TensorFlow Serving, TensorFlow JavaScript.

You can also find usage examples in other projects:

  * [Vision Transformer models converted to TensorFlow](https://git.io/Jz8pj) and integrated with TensorFlow Hub.

# Generating TensorFlow SavedModel

The functions generated by `jax2tf.convert` are standard TensorFlow functions
and you can save them in a SavedModel using standard TensorFlow code, as shown
in the [jax2tf documentation](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#usage-saved-model).
This decoupling of jax2tf from SavedModel is important, because it **allows
the user to have full control over what metadata is saved in the SavedModel**.

As an example, we provide the function `convert_and_save_model`
(see [saved_model_lib.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/saved_model_lib.py).)
For serious uses, you will probably want to copy and expand this
function as needed.

The `convert_and_save_model` example function allows you to control
which model parameters you want to save separately as variables,
as opposed to having them embedded as constants in the function.
This is useful for two reasons:
the parameters could be very large and exceed the limits of the
GraphDef part of the SavedModel, or you may want to fine-tune the
model and change the value of the parameters.

To achieve this you should extract a pair from your trained model:

  * `predict_fn`: a two-argument function with signature:
   `(params: Parameters, inputs: Inputs) -> Outputs`.
   Both arguments must be a `numpy.ndarray` or
   (nested) tuple/list/dictionaries thereof. In particular,
   it is important for models that take multiple inputs to be packaged
   as a function with just two arguments, e.g., by taking their inputs
   as a single tuple/list/dictionary.
  * `params`: of type `Parameters`, the model parameters, to be used as the
  first input for `predict_fn`, and to be saved as the variables of the
  SavedModel. This must also be a `numpy.ndarray` or
  (nested) tuple/list/dictionary thereof.

## Usage: saved models from Flax

If you are using Flax, then the recipe to obtain this pair is as follows:

  ```python
  class MyModel(nn.Module):
     ...

  model = MyModel(*config_args, **kwargs)  # Construct the model
  optimizer = ...  # Train the model
  params = optimizer.target  # These are the parameters
  predict_fn = lambda params, input: model.apply({"params": params}, input)
  ```

You can see in
[mnist_lib.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/mnist_lib.py)
how this can be done for two
implementations of MNIST, one using pure JAX (`PureJaxMNIST`) and a CNN
one using Flax (`FlaxMNIST`). Other Flax models can be arranged similarly,
and the same strategy should work for other neural-network libraries for JAX.

If your Flax model takes multiple inputs, then you need to change the last
line in the example above to:

  ```python
  predict_fn = lambda params, input: model.apply({"params": params}, *input)
  ```

(Notice that `predict_fn` takes a single tuple input, and expands it to
pass it to `model.apply` as multiple inputs.)

You can control which parameters you want to save as variables, and 
which you want to embed in the computation graph. For example, to 
embed all parameters in the graph:

  ```python
  params = ()
  predict_fn = lambda _, input: model.apply({"params": optimizer.target}, input)
  ```

(The MNIST Flax examples from
[mnist_lib.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/mnist_lib.py)
normally has a GraphDef of 150k and a variables section of 3Mb. If we embed the
parameters as constants in the GraphDef as shown above, the variables section
becomes empty and the GraphDef becomes 13Mb. This embedding may allow
the compiler to generate faster code by constant-folding some computations.)

## Usage: saved models from Haiku

If you are using Haiku, then the recipe is along these lines:

  ```python
  model_fn = ...define your Haiku model...
  net = hk.transform(model_fn)  # get a pair of init and apply functions
  params = ... train your model starting from net.init() ...
  predict_fn = hk.without_apply_rng(net).apply
  ```

# Preparing the model for jax2tf

Once you have the model in this form, you can use the
`saved_model_lib.save_model` function from
[saved_model_lib.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/saved_model_lib.py)
to generate the SavedModel. There is very little in that function that is
specific to jax2tf. The goal of jax2tf is to convert JAX functions
into functions that behave as if they had been written with TensorFlow.
Therefore, if you are familiar with how to generate SavedModel, you can most
likely just use your own code for this.

The file
[saved_model_main.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/saved_model_main.py)
is an executable that shows how to perform the
following sequence of steps:

   * train an MNIST model, and obtain a pair of an inference function and the
   parameters.
   * convert the inference function with jax2tf, for one or more batch sizes.
   * save a SavedModel and dump its contents.
   * reload the SavedModel and run it with TensorFlow to test that the inference
   function produces the same results as the JAX inference function.
   * optionally plot images with the training digits and the inference results.


There are a number of flags to select the Flax model (`--model=mnist_flax`),
to skip the training and just test a previously loaded
SavedModel (`--nogenerate_model`), to choose the saving path, etc.
The default saving location is `/tmp/jax2tf/saved_models/`.


By default, this example will convert the inference function for three separate
batch sizes: 1, 16, 128. You can see this in the dumped SavedModel.

# Reusing models with TensorFlow Hub and TensorFlow Keras

The SavedModel produced by the example in `saved_model_main.py` already
implements the [reusable saved models interface](https://www.tensorflow.org/hub/reusable_saved_models).
The executable
[keras_reuse_main.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/keras_reuse_main.py)
extends
[saved_model_main.py](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/saved_model_main.py)
with code
to include a jax2tf SavedModel into a larger TensorFlow Keras model.

In order to demonstrate reuse, we have added a keyword argument `with_classifier`
to the MNIST models. When this flag is set to False, the model will not include
the last classification layer. The SavedModel will compute just the logits.

We show in
`keras_reuse_main.py` how to use `hub.KerasLayer` to turn the
SavedModel into a Keras layer, to which we add a Keras classification layer,
that is then trained in TensorFlow.

For ease of training, we show this with the high-level Keras API, but it should
also be possible in low-level TensorFlow (using the restored_features under
a `tf.GradientTape` and doing gradient updates manually).

All the flags of `saved_model_main.py` also apply to `keras_reuse_main.py`.
In particular, you can select the Flax MNIST model: `--model=mnist_flax`.

# Using jax2tf with TensorFlow serving

It is also possible to use jax2tf-generated SavedModel with TensorFlow serving.
At the moment, the open-source TensorFlow model server is missing XLA support,
but the Google version can be used, as shown in the
[serving examples](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/serving/README.md).

# Using jax2tf with TensorFlow JavaScript

A jax2tf-generated SavedModel can also be converted to a format usable with
TensorFlow.js, by using the appropriate converters from SavedModel. 
At the moment, these converters may reject some jax2tf-generated SavedModels due to
some ops not yet implemented in the converters. As a partial workaround, one
can pass the `enable_xla=False` parameter to `jax2tf.convert` to direct 
`jax2tf` to avoid problematic ops. This will increase the coverage, and in fact
most, but not all, Flax examples can be converted this way. 

Check out
[Quickdraw TensorFlow.js example](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/examples/tf_js/quickdraw/README.md).
