Merge pull request #21803 from gnecula:export_typing

PiperOrigin-RevId: 642251107

Author: jax authors

jax/_src/export/_export.py

@@ -286,9 +286,9 @@ class Exported:
   """
   fun_name: str
   in_tree: tree_util.PyTreeDef
-  in_avals: tuple[core.AbstractValue, ...]
+  in_avals: tuple[core.ShapedArray, ...]
   out_tree: tree_util.PyTreeDef
-  out_avals: tuple[core.AbstractValue, ...]
+  out_avals: tuple[core.ShapedArray, ...]
 
   in_shardings_hlo: tuple[HloSharding | None, ...]
   out_shardings_hlo: tuple[HloSharding | None, ...]
@@ -1257,8 +1257,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
+    if not isinstance(actual_aval, core.ShapedArray):
+      raise ValueError(f"Expected ShapedArray but got: {actual_aval}")
     def pp_arg_dim(dim_idx: int | None) -> str:
       return shape_poly.pretty_print_dimension_descriptor(exported.in_tree,
                                                           arg_idx, dim_idx)
@@ -1302,10 +1302,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,  # type: ignore
+      core.ShapedArray(core.evaluate_shape(out_aval.shape, exported_dim_vars,
                                            *exported_dim_values),
-                       dtype=out_aval.dtype, weak_type=out_aval.weak_type,  # type: ignore
-                       named_shape=out_aval.named_shape)  # type: ignore
+                       dtype=out_aval.dtype, weak_type=out_aval.weak_type,
+                       named_shape=out_aval.named_shape)
       for out_aval in exported.out_avals)
   return out_avals, set(exported.ordered_effects + exported.unordered_effects)
 

jax/_src/export/serialization.fbs

@@ -38,7 +38,7 @@ table PyTreeDef {
 
 enum AbstractValueKind: byte {
   shapedArray = 0,
-  abstractToken = 1,
+  abstractToken = 1,  // unused
 }
 
 enum DType: byte {

jax/_src/export/serialization.py

@@ -331,28 +331,24 @@ def _deserialize_pytreedef_to_pytree(p: ser_flatbuf.PyTreeDef):
 
 
 def _serialize_aval(
-    builder: flatbuffers.Builder, aval: core.AbstractValue
+    builder: flatbuffers.Builder, aval: core.ShapedArray
 ) -> int:
-  aval_type = type(aval)
-  if aval_type is core.ShapedArray:
-    aval_kind = ser_flatbuf.AbstractValueKind.shapedArray
-    shape_offsets = [builder.CreateString(str(d)) for d in aval.shape]
-    ser_flatbuf.AbstractValueStartShapeVector(builder, len(aval.shape))
-    for d in reversed(shape_offsets):
-      builder.PrependUOffsetTRelative(d)
-    shape_vector_offset = builder.EndVector()
-
-    ser_flatbuf.AbstractValueStart(builder)
-    ser_flatbuf.AbstractValueAddKind(builder, aval_kind)
-    ser_flatbuf.AbstractValueAddShape(builder, shape_vector_offset)
-    ser_flatbuf.AbstractValueAddDtype(builder, _dtype_to_dtype_kind[aval.dtype])
-    return ser_flatbuf.AbstractValueEnd(builder)
-  else:
-    raise NotImplementedError(f"serializing AbstractValue: {aval}")
+  aval_kind = ser_flatbuf.AbstractValueKind.shapedArray
+  shape_offsets = [builder.CreateString(str(d)) for d in aval.shape]
+  ser_flatbuf.AbstractValueStartShapeVector(builder, len(aval.shape))
+  for d in reversed(shape_offsets):
+    builder.PrependUOffsetTRelative(d)
+  shape_vector_offset = builder.EndVector()
+
+  ser_flatbuf.AbstractValueStart(builder)
+  ser_flatbuf.AbstractValueAddKind(builder, aval_kind)
+  ser_flatbuf.AbstractValueAddShape(builder, shape_vector_offset)
+  ser_flatbuf.AbstractValueAddDtype(builder, _dtype_to_dtype_kind[aval.dtype])
+  return ser_flatbuf.AbstractValueEnd(builder)
 
 
 def _deserialize_aval(aval: ser_flatbuf.AbstractValue,
-                      scope) -> core.AbstractValue:
+                      scope) -> core.ShapedArray:
   aval_kind = aval.Kind()
   if aval_kind == ser_flatbuf.AbstractValueKind.shapedArray:
     dtype = _dtype_kind_to_dtype[aval.Dtype()]

jax/_src/export/shape_poly.py

@@ -1717,10 +1717,10 @@ def _dimension_size_lowering_rule(ctx, arg, *, dimension):
 mlir.register_lowering(dimension_size_p, _dimension_size_lowering_rule)
 
 
-def all_dim_vars(args_avals: Sequence[core.AbstractValue]) -> Sequence[str]:
+def all_dim_vars(args_avals: Sequence[core.ShapedArray]) -> Sequence[str]:
   dim_vars: set[str] = set()
   for a in args_avals:
-    for d in a.shape:  # type: ignore[attribute-error,unused-ignore]
+    for d in a.shape:
       if is_symbolic_dim(d):
         dim_vars = dim_vars.union(d._get_vars())
   return sorted(dim_vars)
@@ -1911,7 +1911,7 @@ def pretty_print_dimension_descriptor(
 
 @util.cache()
 def solve_dim_vars(
-    args_avals: Sequence[core.AbstractValue],
+    args_avals: Sequence[core.ShapedArray],
     args_kwargs_tree: tree_util.PyTreeDef,
     ) -> tuple[DimVarEnv, ShapeConstraints, Sequence[tuple[str, int, int]]]:
   """Solves dimension variables in a called function's avals in terms of actual argument shapes.
@@ -1956,12 +1956,12 @@ def solve_dim_vars(
   # tuples with argument name and its polymorphic shape ('args[0]', '(a, a + b'))
   polymorphic_shape_specs: list[tuple[str, str]] = []
   for arg_idx, aval in enumerate(args_avals):
-    if all(not is_symbolic_dim(d) for d in aval.shape):  # type: ignore[attribute-error,unused-ignore]
+    if all(not is_symbolic_dim(d) for d in aval.shape):
       continue
     polymorphic_shape_specs.append(
       (pretty_print_dimension_descriptor(args_kwargs_tree, arg_idx, None),
-       str(aval.shape)))  # type: ignore[attribute-error,unused-ignore]
-    for dim_idx, aval_d in enumerate(aval.shape):  # type: ignore[attribute-error,unused-ignore]
+       str(aval.shape)))
+    for dim_idx, aval_d in enumerate(aval.shape):
       if is_symbolic_dim(aval_d):
         synth_dim_var = pretty_print_dimension_descriptor(args_kwargs_tree,
                                                           arg_idx, dim_idx)
@@ -1976,7 +1976,7 @@ def solve_dim_vars(
 
 
 def compute_dim_vars_from_arg_shapes(
-    args_avals: Sequence[core.AbstractValue],
+    args_avals: Sequence[core.ShapedArray],
     *actual_args: jax.Array,
     args_kwargs_tree: tree_util.PyTreeDef) -> Sequence[jax.Array]:
   """Computes values of dimension variables to unify args_avals with actual arguments.

tests/export_test.py

@@ -1489,10 +1489,10 @@ def f_outer(x):
       self.assertAllClose(2. * 2. * x + 10. + 4. * 2. * x, res)
 
   @jtu.parameterized_filterable(
-    kwargs=[
-      dict(v=v)
-      for v in range(export.minimum_supported_serialization_version,
-                     export.maximum_supported_serialization_version + 1)])
+      kwargs=[
+          dict(v=v)
+          for v in range(export.minimum_supported_serialization_version,
+                         export.maximum_supported_serialization_version + 1)])
   def test_ordered_effects_poly(self, *, v: int):
     with config.jax_serialization_version(v):
       logging.info(