Merge pull request #21763 from gnecula:export_api

PiperOrigin-RevId: 641959833

Author: jax authors

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