Plans with jax.experimental.export

[dionhaefner]

Hey all,

I noticed that jax.experimental.export landed in the latest JAX release. I take it that this module is meant as a way to serialize and load JAX graphs, as an alternative to using jax2tf to generate a tensorflow SavedModel?

I'm currently evaluating different ways to deploy JAX models, and a native way to save and load jitted inference functions would be a strong contender. Given that there are no docs for the module yet, could you give us a glimpse of what's on the horizon?

1. Are there any general restrictions to keep in mind? (in regards to supported ops, vmap, gradients, platforms, sharding etc.)
2. Is it reasonable to expect that this could be used to save and load non-trivial flax modules within a couple of months from now?
3. Any plans regarding forward compatibility with other JAX versions?

[gnecula]

Hi,

We are preparing the documentation for the export module, but you are right that it is meant to be an alternative to jax2tf and SavedModel. In fact, the logic of the module is to a large extent extracted from jax2tf, minus the TF parts, plus a few features that TF does not support, such as effects. You can get a good sense of the features supported and the limitations if you look at the jax2tf documentation (the "native serialization" parts).

We intend to support all JAX programs, although at the moment pmap and xmap are not supported. vmap is supported in the sense that you can export a program that uses vmap. The same is true for other transformations. There are however limitations of what you can do with an exported program: you can call it, and you can compute higher-order gradients. But you cannot apply other transformations on it. We may expand on this at some point.

This module should already be ready to use for saving and loading Flax models, and other JAX programs.

The serialization part of export is newer, and it may change. However, we will support old serialized exported modules for at least 6 months, meaning that you will be able to deserialized and call an exported module with versions of JAX newer by at most 6 months than the version at which the code was exported.

[dionhaefner]

That sounds awesome, thanks!

[angela-ko]

Since the documentation is still in progress, are there any example usages of the export module I could reference in the mean time? I don't seem to see an examples / tests file in the export directory. Thanks!

[gnecula]

There are tests in the tests/export_test.py file.

[stellaraccident]

Do you guys have a test suite of exported models yet? We'd like to support this better and it is fairly easy to run a project that has a good test suite of model variants (based on experience with onnx and torch). It's harder if the backend team has to synthesize that because it requires pulling from the more limited pool of folks who are competent to be putting together a Jax model garden.

[gnecula]

We do not have a test suite of exported models. There are other repos like Flax and MaxText that have model examples, although not in exported form.

For backend testing, the best we have are the JAX tests. They have really good coverage. There is even a way to generate several thousand test harnesses that should have good coverage for JAX primitives. See how test_harness is used in export_harnesses_multi_platform_test.py