[export] Move the export implementation to jax._src.export.

This is part of the work to move the export APIs out
of jax.experimental. For now, the way to use this
implementation is still through `jax.experimental.export`.

Had to add a few "#type ignore" to the _export.py because
previously the file was exempt from internal pytype.
Will try to fix these in a later PR.

PiperOrigin-RevId: 641688200

Author: gnecula

jax/BUILD

@@ -321,7 +321,7 @@ py_library_providing_imports_info(
         ":xla",
         ":xla_bridge",
         "//jax/_src/lib",
-    ] + py_deps("numpy") + py_deps("scipy") + py_deps("opt_einsum") + jax_extra_deps,
+    ] + py_deps("numpy") + py_deps("scipy") + py_deps("opt_einsum") + py_deps("flatbuffers") + jax_extra_deps,
 )
 
 pytype_strict_library(

jax/_src/export/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

jax/experimental/export/_export.py renamed to jax/_src/export/_export.py

@@ -460,7 +460,7 @@ def do_export(*args_specs, **kwargs_specs) -> Exported:
       actual_lowering_platforms = (default_lowering_platform(),)
 
     # TODO: move to `lower`
-    symbolic_scope: tuple[shape_poly.SymbolicScope, tree_util.KeyPath] | None = None
+    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"):
@@ -476,7 +476,7 @@ def do_export(*args_specs, **kwargs_specs) -> Exported:
               other_descr=shape_poly.args_kwargs_path_to_str(k_path))
 
     if has_trace:
-      traced = wrapped_fun_jax.trace(
+      traced = wrapped_fun_jax.trace(  # type: ignore
           *args_specs, **kwargs_specs,
           _experimental_lowering_parameters=mlir.LoweringParameters(
             platforms=actual_lowering_platforms,
@@ -547,7 +547,7 @@ def _export_lowered(
   elif "shards" in lowering.compile_args:  # for PmapComputation
     out_avals_flat = lowering.compile_args["shards"].out_sharded_avals
   else:
-    out_avals_flat = lowered.compile_args["out_avals"]
+    out_avals_flat = lowered.compile_args["out_avals"]  # type: ignore
 
   # Log and then check the module.
   if logging.vlog_is_on(3):
@@ -612,7 +612,7 @@ def _get_exported_vjp(exp_primal: Exported) -> Exported:
       in_shardings_hlo=in_shardings,
       out_shardings_hlo=out_shardings,
       nr_devices=nr_devices,
-      lowering_platforms=lowering._platforms,
+      lowering_platforms=lowering._platforms,  # type: ignore
       ordered_effects=ordered_effects,
       unordered_effects=unordered_effects,
       disabled_safety_checks=tuple(disabled_checks),
@@ -641,7 +641,7 @@ def _module_to_bytecode(module: ir.Module) -> bytes:
   # and still have the payloads produced by `serialize_portable_artifact`
   # compatible with potential consumers from the past.
   target_version = hlo.get_minimum_version()
-  module_serialized = xla_client._xla.mlir.serialize_portable_artifact(
+  module_serialized = xla_client._xla.mlir.serialize_portable_artifact(  # type: ignore
       mlir_str, target_version)
   return module_serialized
 
@@ -688,8 +688,8 @@ def _wrap_main_func(
     def is_token(typ, attrs):
       return (typ == mlir.token_type()[0])
 
-    orig_input_types = orig_main.type.inputs
-    arg_attrs = list(ir.ArrayAttr(orig_main.arg_attrs))
+    orig_input_types = orig_main.type.inputs  # type: ignore
+    arg_attrs = list(ir.ArrayAttr(orig_main.arg_attrs))  # type: ignore
     # The order of args: platform_index_arg, dim args, token args, array args.
     nr_platform_index_args = 1 if has_platform_index_argument else 0
     nr_dim_args = len(dim_vars)
@@ -711,8 +711,8 @@ def is_token(typ, attrs):
       orig_input_types, [nr_platform_index_args, nr_dim_args, nr_token_args])
 
     # The order of results: tokens, array results
-    orig_output_types = orig_main.type.results
-    result_attrs = list(ir.ArrayAttr(orig_main.result_attrs))
+    orig_output_types = orig_main.type.results  # type: ignore
+    result_attrs = list(ir.ArrayAttr(orig_main.result_attrs))  # type: ignore
     token_result_idxs = [i for i, (typ, attrs) in enumerate(zip(orig_output_types,
                                                                 result_attrs))
                          if is_token(typ, attrs)]
@@ -1138,6 +1138,8 @@ def _call_exported_abstract_eval(
   assert len(in_avals) == len(exported.in_avals)  # since the pytrees have the same structure
   # Check that the expected shapes match the actual ones
   for arg_idx, (exp_aval, actual_aval) in enumerate(zip(exported.in_avals, in_avals)):
+    exp_aval: core.ShapedArray = exp_aval  # type: ignore
+    actual_aval: core.ShapedArray = actual_aval  # type: ignore
     def pp_arg_dim(dim_idx: int | None) -> str:
       return shape_poly.pretty_print_dimension_descriptor(exported.in_tree,
                                                           arg_idx, dim_idx)
@@ -1181,10 +1183,10 @@ def pp_arg_dim(dim_idx: int | None) -> str:
   exported_dim_values = [synthetic_eval.evaluate(solution[var])
                          for var in exported_dim_vars]
   out_avals = tuple(
-      core.ShapedArray(core.evaluate_shape(out_aval.shape, exported_dim_vars,
+      core.ShapedArray(core.evaluate_shape(out_aval.shape, exported_dim_vars,  # type: ignore
                                            *exported_dim_values),
-                       dtype=out_aval.dtype, weak_type=out_aval.weak_type,
-                       named_shape=out_aval.named_shape)
+                       dtype=out_aval.dtype, weak_type=out_aval.weak_type,  # type: ignore
+                       named_shape=out_aval.named_shape)  # type: ignore
       for out_aval in exported.out_avals)
   return out_avals, set(exported.ordered_effects + exported.unordered_effects)
 

jax/experimental/export/serialization.fbs renamed to jax/_src/export/serialization.fbs

@@ -20,7 +20,7 @@
 //  3.  Add back the licence comment at the start
 //
 
-namespace jax.experimental.export.serialization;
+namespace jax.export.serialization;
 
 enum PyTreeDefKind: byte {
   leaf = 0,

jax/experimental/export/_serialization.py renamed to jax/_src/export/serialization.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Serialization and deserialization of export.Exported
+# Serialization and deserialization of _export.Exported
 
 from __future__ import annotations
 
@@ -31,10 +31,10 @@
 from jax._src import dtypes
 from jax._src import effects
 from jax._src import tree_util
+from jax._src.export import serialization_generated as ser_flatbuf
+from jax._src.export import _export
+from jax._src.export import shape_poly
 from jax._src.lib import xla_client
-from jax.experimental.export import serialization_generated as ser_flatbuf
-from jax.experimental.export import _export
-from jax.experimental import export
 
 import numpy as np
 
@@ -47,7 +47,7 @@
 # Version 2, Dec 16th, 2023, adds the f0 dtype.
 _SERIALIZATION_VERSION = 2
 
-def serialize(exp: export.Exported, vjp_order: int = 0) -> bytearray:
+def serialize(exp: _export.Exported, vjp_order: int = 0) -> bytearray:
   """Serialize an Exported.
 
   Args:
@@ -63,14 +63,14 @@ def serialize(exp: export.Exported, vjp_order: int = 0) -> bytearray:
   return builder.Output()
 
 
-def deserialize(ser: bytearray) -> export.Exported:
+def deserialize(ser: bytearray) -> _export.Exported:
   """Deserialize an Exported."""
   exp = ser_flatbuf.Exported.GetRootAsExported(ser)
   return _deserialize_exported(exp)
 
 
 def _serialize_exported(
-    builder: flatbuffers.Builder, exp: export.Exported, vjp_order: int
+    builder: flatbuffers.Builder, exp: _export.Exported, vjp_order: int
 ) -> int:
   # Serialize bottom-up
   fun_name = builder.CreateString(exp.fun_name)
@@ -150,7 +150,7 @@ def _serialize_array(
   return builder.EndVector()
 
 
-def _deserialize_exported(exp: ser_flatbuf.Exported) -> export.Exported:
+def _deserialize_exported(exp: ser_flatbuf.Exported) -> _export.Exported:
   serialization_version = exp.SerializationVersion()
   if serialization_version != _SERIALIZATION_VERSION:
     raise NotImplementedError(
@@ -161,7 +161,7 @@ def _deserialize_exported(exp: ser_flatbuf.Exported) -> export.Exported:
   _, in_tree = tree_util.tree_flatten(
       _deserialize_pytreedef_to_pytree(exp.InTree())
   )
-  scope = export.SymbolicScope(())  # TODO: serialize the constraints
+  scope = shape_poly.SymbolicScope(())  # TODO: serialize the constraints
   deser_aval = partial(_deserialize_aval, scope=scope)
   in_avals = _deserialize_tuple(
       exp.InAvalsLength, exp.InAvals, deser_aval
@@ -205,7 +205,7 @@ def _deserialize_exported(exp: ser_flatbuf.Exported) -> export.Exported:
   if vjp := exp.Vjp():
     _get_vjp = lambda _: _deserialize_exported(vjp)
 
-  return export.Exported(
+  return _export.Exported(
       fun_name=fun_name,
       in_tree=in_tree,
       in_avals=in_avals,
@@ -356,7 +356,7 @@ def _deserialize_aval(aval: ser_flatbuf.AbstractValue,
   aval_kind = aval.Kind()
   if aval_kind == ser_flatbuf.AbstractValueKind.shapedArray:
     dtype = _dtype_kind_to_dtype[aval.Dtype()]
-    shape = export.symbolic_shape(
+    shape = shape_poly.symbolic_shape(
         ",".join(
             aval.Shape(i).decode("utf-8") for i in range(aval.ShapeLength())
         ),
@@ -445,16 +445,16 @@ def _deserialize_effect(eff: ser_flatbuf.Effect) -> core.Effect:
 
 
 def _serialize_disabled_safety_check(
-    builder: flatbuffers.Builder, check: export.DisabledSafetyCheck
+    builder: flatbuffers.Builder, check: _export.DisabledSafetyCheck
 ) -> int:
   custom_call_target_str = check.is_custom_call()
   custom_call_target = None
   if custom_call_target_str is not None:
     kind = ser_flatbuf.DisabledSafetyCheckKind.custom_call
     custom_call_target = builder.CreateString(custom_call_target_str)
-  elif check == export.DisabledSafetyCheck.platform():
+  elif check == _export.DisabledSafetyCheck.platform():
     kind = ser_flatbuf.DisabledSafetyCheckKind.platform
-  elif check == export.DisabledSafetyCheck.shape_assertions():
+  elif check == _export.DisabledSafetyCheck.shape_assertions():
     kind = ser_flatbuf.DisabledSafetyCheckKind.shape_assertions
   else:
     raise NotImplementedError(f"serializing DisabledSafetyCheck: {check}")
@@ -470,14 +470,14 @@ def _serialize_disabled_safety_check(
 
 def _deserialize_disabled_safety_check(
     sc: ser_flatbuf.DisabledSafetyCheck,
-) -> export.DisabledSafetyCheck:
+) -> _export.DisabledSafetyCheck:
   kind = sc.Kind()
   if kind == ser_flatbuf.DisabledSafetyCheckKind.custom_call:
-    return export.DisabledSafetyCheck.custom_call(
+    return _export.DisabledSafetyCheck.custom_call(
         sc.CustomCallTarget().decode("utf-8")
     )
   if kind == ser_flatbuf.DisabledSafetyCheckKind.platform:
-    return export.DisabledSafetyCheck.platform()
+    return _export.DisabledSafetyCheck.platform()
   if kind == ser_flatbuf.DisabledSafetyCheckKind.shape_assertions:
-    return export.DisabledSafetyCheck.shape_assertions()
+    return _export.DisabledSafetyCheck.shape_assertions()
   assert False, kind

jax/experimental/export/serialization_generated.py renamed to jax/_src/export/serialization_generated.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pytype: skip-file
 # automatically generated by the FlatBuffers compiler, do not modify
 
 # namespace: serialization

jax/experimental/export/BUILD

@@ -31,9 +31,6 @@ py_library(
     name = "export",
     srcs = [
         "__init__.py",
-        "_export.py",
-        "_serialization.py",
-        "serialization_generated.py",
     ],
     srcs_version = "PY3",
     # TODO: b/255503696: enable pytype

jax/experimental/export/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from jax.experimental.export._export import (
+from jax._src.export._export import (
     minimum_supported_serialization_version,
     maximum_supported_serialization_version,
     Exported,
@@ -29,7 +29,7 @@
     symbolic_args_specs,
     SymbolicScope,
 )
-from jax.experimental.export._serialization import (
+from jax._src.export.serialization import (
     serialize,
     deserialize,
 )

jax/experimental/jax2tf/jax2tf.py

@@ -37,7 +37,6 @@
 from jax import tree_util
 from jax import sharding
 from jax.experimental import export
-from jax.experimental.export import _export
 from jax.experimental.jax2tf import impl_no_xla
 from jax.interpreters import xla
 
@@ -60,6 +59,7 @@
 from jax._src import source_info_util
 from jax._src import util
 from jax._src import shard_alike
+from jax._src.export import _export
 from jax._src.export import shape_poly
 from jax._src.interpreters import ad
 from jax._src.interpreters import mlir

tests/export_back_compat_test.py

@@ -27,7 +27,7 @@
 
 import jax
 from jax import lax
-from jax.experimental.export import _export
+from jax._src.export import _export
 
 from jax._src.internal_test_util import export_back_compat_test_util as bctu