[export] Create the jax.export module APIs.

The functionality comes from the jax.experimental.export
module, which will be deprecated.

The following APIs are introduced:

```
  from jax import export
  def f(...): ...
  ex: export.Exported = export.export(jax.jit(f))(*args, **kwargs)

  blob: bytearray = ex.serialize()
  rehydrated: export.Export = export.deserialize(blob)

  def caller(...):
     ... rehydrated.call(*args, **kwargs)
```

Module documentation will follow shortly.
There are no changes for now in the jax.experimental.export
APIs.

Most of the changes in this PR are in tests due to some differences
in the new jax.export APIs compared to jax.experimental.export:

  * Instead of `jax.experimental.export.call(exp)` we now write
    `exp.call`
  * The `jax.experimental.export.export` allowed the function
    argument to be any Python callable and it would wrap it with
    a `jax.jit`. This is not supported anymore by export, and instead
    the user must use `jax.jit`.

Author: gnecula

benchmarks/shape_poly_benchmark.py

@@ -18,7 +18,7 @@
 import jax
 from jax import core
 from jax._src.numpy import lax_numpy
-from jax.experimental import export
+from jax import export
 
 jax.config.parse_flags_with_absl()
 

jax/_src/export/_export.py

@@ -333,10 +333,10 @@ def in_shardings_jax(
     `jax.device_put`.
 
     Example usage:
-    >>> from jax.experimental import export
+    >>> from jax import export
     >>> exp_mesh = sharding.Mesh(jax.devices(), ("a",))
     >>> exp = export.export(jax.jit(lambda x: jax.numpy.add(x, x),
-    ...     in_shardings=sharding.NamedSharding(exp_mesh, sharding.PartitionSpec("a")))
+    ...                             in_shardings=sharding.NamedSharding(exp_mesh, sharding.PartitionSpec("a")))
     ...     )(np.arange(jax.device_count()))
     >>> exp.in_shardings_hlo
     ({devices=[8]<=[8]},)
@@ -347,7 +347,7 @@ def in_shardings_jax(
     # Put the args and kwargs on the appropriate devices
     >>> run_arg = jax.device_put(np.arange(jax.device_count()),
     ...     exp.in_shardings_jax(run_mesh)[0])
-    >>> res = export.call(exp)(run_arg)
+    >>> res = exp.call(run_arg)
     >>> res.addressable_shards
     [Shard(device=CpuDevice(id=7), index=(slice(0, 1, None),), replica_id=0, data=[0]),
      Shard(device=CpuDevice(id=6), index=(slice(1, 2, None),), replica_id=0, data=[2]),
@@ -372,19 +372,53 @@ def out_shardings_jax(
                  for s in self.out_shardings_hlo)
 
   def has_vjp(self) -> bool:
+    """Returns if this Exported supports VJP."""
     return self._get_vjp is not None
 
   def vjp(self) -> Exported:
     """Gets the exported VJP.
 
     Returns None if not available, which can happen if the Exported has been
-    loaded from an external format, without a VJP."""
+    loaded from an external format without a VJP.
+    """
     if self._get_vjp is None:
       raise ValueError("No VJP is available")
     return self._get_vjp(self)
 
+  def serialize(self,
+                vjp_order: int = 0) -> bytearray:
+    """Serializes an Exported.
+
+    Args:
+      vjp_order: The maximum vjp order to include. E.g., the value 2 means that we
+        serialize the primal functions and two orders of the `vjp` function. This
+        should allow 2nd order reverse mode differentiation of the deserialized
+        function. i.e., `jax.grad(jax.grad(f)).`
+    """
+    # Lazy load the serialization module, since flatbuffers is an optional
+    # dependency.
+    from jax._src.export.serialization import serialize
+    return serialize(self, vjp_order=vjp_order)
+
+  def call(self, *args, **kwargs):
+    return call_exported(self)(*args, **kwargs)
+
+
+def deserialize(blob: bytearray) -> Exported:
+  """Deserializes an Exported.
+
+  Args:
+    blob: a bytearray obtained from `Exported.serialize`.
+  """
+  # Lazy load the serialization module, since flatbuffers is an optional
+  # dependency.
+  from jax._src.export.serialization import deserialize
+  return deserialize(blob)
+
 
 def default_lowering_platform() -> str:
+  """Retrieves the default lowering platform for the exporting machine.
+  """
   # Canonicalize to turn 'gpu' into 'cuda' or 'rocm'
   return xb.canonicalize_platform(jax.default_backend())
 
@@ -411,14 +445,20 @@ def args_specs(
   return shape_poly.symbolic_args_specs(args, polymorphic_shapes)
 
 
-def export(fun_jax: Callable,
-           *,
-           lowering_platforms: Sequence[str] | None = None,
-           disabled_checks: Sequence[DisabledSafetyCheck] = (),
-           _device_assignment_for_internal_jax2tf_use_only = None,
-           ) -> Callable[..., Exported]:
+# TODO(necula): remove this once we remove jax.experimental.export.
+def export_back_compat(
+    fun_jax: Callable,
+    *,
+    lowering_platforms: Sequence[str] | None = None,
+    disabled_checks: Sequence[DisabledSafetyCheck] = (),
+    _device_assignment_for_internal_jax2tf_use_only = None,
+    ) -> Callable[..., Exported]:
   """Exports native serialization for a JAX function.
 
+  Note: this function exists only for internal usage by jax2tf and for
+    backwards compatibility with jax.experimental.export. Use
+    `jax.export` instead.
+
   Args:
     fun_jax: the function to lower and serialize.
     lowering_platforms:
@@ -498,6 +538,85 @@ def do_export(*args_specs, **kwargs_specs) -> Exported:
         _device_assignment_for_internal_jax2tf_use_only=_device_assignment_for_internal_jax2tf_use_only)
   return do_export
 
+def export(
+    fun_jit: stages.Wrapped,
+    *,
+    lowering_platforms: Sequence[str] | None = None,
+    disabled_checks: Sequence[DisabledSafetyCheck] = (),
+    ) -> Callable[..., Exported]:
+  """Exports a JAX function for persistent serialization.
+
+  Args:
+    fun_jit: the function to export. Should be the result of `jit`.
+    lowering_platforms:
+        Optional sequence containing a subset of 'tpu', 'cpu',
+        'cuda', 'rocm'. If more than one platform is specified, then
+        the lowered code takes an argument specifying the platform.
+        If None, then use the default JAX backend.
+        The calling convention for multiple platforms is explained in the
+        `jax_export.Exported` docstring.
+    disabled_checks: the safety checks to disable. See docstring
+        of `DisabledSafetyCheck`.
+
+  Returns: a function that takes args and kwargs pytrees of jax.ShapeDtypeStruct,
+      or values with `.shape` and `.dtype` attributes, and returns an
+      `Exported`.
+
+  Usage:
+      >>> from jax import export
+      >>> exported: export.Exported = export.export(jnp.sin)(
+      ...     np.arange(4, dtype=np.float32))
+
+      # You can inspect the Exported object
+      >>> exported.in_avals
+      (ShapedArray(float32[4]),)
+      >>> blob: bytearray = exported.serialize()
+
+      # The serialized bytes are safe to use in a separate process
+      >>> rehydrated: export.Exported = export.deserialize(blob)
+      >>> rehydrated.fun_name
+      'sin'
+      >>> rehydrated.call(np.array([.1, .2, .3, .4], dtype=np.float32))
+      Array([0.09983342, 0.19866933, 0.29552022, 0.38941833], dtype=float32)
+  """
+  if not isinstance(fun_jit, stages.Wrapped):
+    raise ValueError(
+        f"Function to be exported must be the result of `jit` but is: {fun_jit}")
+  if lowering_platforms is not None:
+    actual_lowering_platforms = tuple(lowering_platforms)
+  else:
+    actual_lowering_platforms = (default_lowering_platform(),)
+
+  def do_export(*args_specs, **kwargs_specs) -> Exported:
+    # TODO: move to `lower`
+    symbolic_scope: tuple[shape_poly.SymbolicScope, tree_util.KeyPath] | None = None  # type: ignore[invalid-annotation,unused-ignore]
+    for k_path, aval in tree_util.tree_flatten_with_path((args_specs, kwargs_specs))[0]:
+      # Static args may have no `shape` attribute.
+      if not hasattr(aval, "shape"):
+        continue
+      for d in aval.shape:
+        if shape_poly.is_symbolic_dim(d):
+          if symbolic_scope is None:
+            symbolic_scope = (d.scope, k_path)
+            continue
+          symbolic_scope[0]._check_same_scope(
+              d, when=f"when exporting {util.fun_name(fun_jit)}",
+              self_descr=f"current (from {shape_poly.args_kwargs_path_to_str(symbolic_scope[1])}) ",
+              other_descr=shape_poly.args_kwargs_path_to_str(k_path))
+
+    traced = fun_jit.trace(  # type: ignore
+        *args_specs, **kwargs_specs,
+        _experimental_lowering_parameters=mlir.LoweringParameters(
+          platforms=actual_lowering_platforms,
+          for_export=True,
+        ))
+    jaxpr, fun_name = traced.jaxpr, traced.fun_name
+    lowered = traced.lower()
+    return _export_lowered(
+        lowered, jaxpr, fun_name,
+        disabled_checks=disabled_checks)
+  return do_export
+
 def _export_lowered(
     lowered: stages.Lowered,
     jaxpr: core.ClosedJaxpr, fun_name: str,
@@ -599,7 +718,7 @@ def _get_exported_vjp(exp_primal: Exported) -> Exported:
                                              device_assignment=device_assignment,
                                              apply_jit=True,
                                              flat_primal_fun=True)
-    return export(fun_vjp_jax,
+    return export(fun_vjp_jax,  # type: ignore[arg-type]
                   lowering_platforms=exp_primal.lowering_platforms,
                   disabled_checks=exp_primal.disabled_safety_checks)(*vjp_in_avals)
 
@@ -816,7 +935,7 @@ def is_token(typ, attrs):
 
 def _check_lowering(lowering) -> None:
   if not isinstance(lowering, pxla.MeshComputation):
-    raise NotImplementedError(f"serialization is supported only for pjit. {lowering}")
+    raise NotImplementedError(f"serialization is supported only for jit. {lowering}")
 
   if lowering.compile_args["host_callbacks"] or lowering.compile_args["keepalive"]:
     raise NotImplementedError("serialization of host_callbacks is not yet implemented")

jax/_src/export/serialization.py

@@ -48,7 +48,7 @@
 _SERIALIZATION_VERSION = 2
 
 def serialize(exp: _export.Exported, vjp_order: int = 0) -> bytearray:
-  """Serialize an Exported.
+  """Serializes an Exported.
 
   Args:
     exp: the Exported to serialize.
@@ -64,7 +64,7 @@ def serialize(exp: _export.Exported, vjp_order: int = 0) -> bytearray:
 
 
 def deserialize(ser: bytearray) -> _export.Exported:
-  """Deserialize an Exported."""
+  """Deserializes an Exported."""
   exp = ser_flatbuf.Exported.GetRootAsExported(ser)
   return _deserialize_exported(exp)
 

jax/_src/internal_test_util/export_back_compat_test_util.py

@@ -86,7 +86,7 @@ def func(...): ...
 
 import jax
 from jax import tree_util
-from jax.experimental import export
+from jax import export
 
 from jax.experimental import pjit
 
@@ -345,4 +345,4 @@ def _get_vjp(_):
       _get_vjp=_get_vjp)
 
       # We use pjit in case there are shardings in the exported module.
-    return pjit.pjit(export.call(exported))(*data.inputs)
+    return pjit.pjit(exported.call)(*data.inputs)

jax/_src/stages.py

@@ -32,7 +32,7 @@
 
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, NamedTuple, Protocol, Union
+from typing import Any, NamedTuple, Protocol, Union, runtime_checkable
 import warnings
 
 import jax
@@ -756,8 +756,9 @@ def cost_analysis(self) -> Any | None:
       return None
 
 
+@runtime_checkable
 class Wrapped(Protocol):
-  """A function ready to be specialized, lowered, and compiled.
+  """A function ready to be traced, lowered, and compiled.
 
   This protocol reflects the output of functions such as
   ``jax.jit``. Calling it results in JIT (just-in-time) lowering,

jax/experimental/export/__init__.py

@@ -17,12 +17,13 @@
     minimum_supported_serialization_version,
     maximum_supported_serialization_version,
     Exported,
-    export,
     call_exported,  # TODO: deprecate
     call,
     DisabledSafetyCheck,
-    default_lowering_platform,
+    default_lowering_platform,  # TODO: deprecate
 )
+from jax._src.export._export import export_back_compat as export
+
 from jax._src.export.shape_poly import (
     is_symbolic_dim,
     symbolic_shape,
@@ -33,4 +34,6 @@
     serialize,
     deserialize,
 )
+# Import only to set the shape poly decision procedure
 from jax._src.export import shape_poly_decision
+del shape_poly_decision

jax/experimental/jax2tf/jax2tf.py

@@ -36,7 +36,7 @@
 from jax import numpy as jnp
 from jax import tree_util
 from jax import sharding
-from jax.experimental import export
+from jax import export
 from jax.experimental.jax2tf import impl_no_xla
 from jax.interpreters import xla
 
@@ -515,7 +515,7 @@ def _restore_context():
 
     self._restore_context = _restore_context
     _exported_device_assignment = [None]
-    self.exported = export.export(
+    self.exported = _export.export_back_compat(
         self.fun_jax,
         lowering_platforms=self.native_serialization_platforms,
         disabled_checks=self.native_serialization_disabled_checks,

jax/experimental/jax2tf/tests/call_tf_test.py

@@ -25,13 +25,13 @@
 import jax
 from jax import dlpack
 from jax import dtypes
+from jax import export
 from jax import lax
 from jax import numpy as jnp
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo
-from jax.experimental import export
 from jax.experimental import jax2tf
 from jax.experimental.jax2tf.tests import tf_test_util
 import numpy as np
@@ -778,7 +778,7 @@ def f_jax(x):
 
     lowering_platforms = ("tpu", "cpu", "cuda")
 
-    exp = export.export(f_jax,
+    exp = export.export(jax.jit(f_jax),
                         lowering_platforms=lowering_platforms)(x)
     for jax_platform in jax_and_tf_platforms:
       with self.subTest(jax_platform):
@@ -787,7 +787,7 @@ def f_jax(x):
         logging.info("Running harness natively on %s", jax_device)
         native_res = f_jax(x_device)
         logging.info("Running exported harness on %s", jax_device)
-        exported_res = export.call(exp)(x_device)
+        exported_res = exp.call(x_device)
         self.assertAllClose(native_res, exported_res)
 
   def test_multi_platform_call_tf_graph(self):

jax/experimental/jax2tf/tests/jax2tf_test.py

@@ -27,6 +27,7 @@
 import jax
 from jax import ad_checkpoint
 from jax import dtypes
+from jax import export
 from jax import lax
 from jax import numpy as jnp
 from jax import sharding
@@ -37,7 +38,6 @@
 from jax._src import test_util as jtu
 from jax._src import xla_bridge as xb
 from jax.experimental import jax2tf
-from jax.experimental import export
 from jax.experimental.jax2tf.tests import tf_test_util
 from jax.experimental.shard_map import shard_map
 from jax.experimental import pjit
@@ -1559,7 +1559,7 @@ def apply_transform(func, transform: str):
       # Run the JAX native version, to check it works, and to fill caches.
       _ = func_to_convert(*args)
       exported = export.export(
-          func_to_convert,
+          (jax.jit(func_to_convert) if not hasattr(func_to_convert, "trace") else func_to_convert),
           lowering_platforms=("tpu",)
       )(*(core.ShapedArray(a.shape, a.dtype) for a in args))
 

jax/experimental/jax2tf/tests/shape_poly_test.py

@@ -32,8 +32,8 @@
 
 import jax
 from jax.experimental import jax2tf
-from jax.experimental import export
 from jax.experimental import pjit
+from jax import export
 from jax import lax
 import jax.numpy as jnp
 from jax import random