Reverts b56dc63

PiperOrigin-RevId: 707501925

Author: Google-ML-Automation

CHANGELOG.md

@@ -12,11 +12,6 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 
 ## Unreleased
 
-* Deprecations
-  * From {mod}`jax.interpreters.xla`, `abstractify` and `pytype_aval_mappings`
-    are now deprecated, having been replaced by symbols of the same name
-    in {mod}`jax.core`. 
-
 ## jax 0.4.38 (Dec 17, 2024)
 
 * Changes:

jax/_src/abstract_arrays.py

@@ -49,6 +49,7 @@ def masked_array_error(*args, **kwargs):
                    "Use arr.filled() to convert the value to a standard numpy array.")
 
 core.pytype_aval_mappings[np.ma.MaskedArray] = masked_array_error
+core.xla_pytype_aval_mappings[np.ma.MaskedArray] = masked_array_error
 
 
 def _make_shaped_array_for_numpy_array(x: np.ndarray) -> ShapedArray:
@@ -57,6 +58,7 @@ def _make_shaped_array_for_numpy_array(x: np.ndarray) -> ShapedArray:
   return ShapedArray(x.shape, dtypes.canonicalize_dtype(dtype))
 
 core.pytype_aval_mappings[np.ndarray] = _make_shaped_array_for_numpy_array
+core.xla_pytype_aval_mappings[np.ndarray] = _make_shaped_array_for_numpy_array
 
 
 def _make_shaped_array_for_numpy_scalar(x: np.generic) -> ShapedArray:
@@ -66,6 +68,7 @@ def _make_shaped_array_for_numpy_scalar(x: np.generic) -> ShapedArray:
 
 for t in numpy_scalar_types:
   core.pytype_aval_mappings[t] = _make_shaped_array_for_numpy_scalar
+  core.xla_pytype_aval_mappings[t] = _make_shaped_array_for_numpy_scalar
 
 core.literalable_types.update(array_types)
 
@@ -78,5 +81,6 @@ def _make_abstract_python_scalar(typ, val):
 
 for t in dtypes.python_scalar_dtypes:
   core.pytype_aval_mappings[t] = partial(_make_abstract_python_scalar, t)
+  core.xla_pytype_aval_mappings[t] = partial(_make_abstract_python_scalar, t)
 
 core.literalable_types.update(dtypes.python_scalar_dtypes.keys())

jax/_src/array.py

@@ -1038,6 +1038,7 @@ def _get_aval_array(self):
 
 api_util._shaped_abstractify_handlers[ArrayImpl] = _get_aval_array
 core.pytype_aval_mappings[ArrayImpl] = _get_aval_array
+core.xla_pytype_aval_mappings[ArrayImpl] = _get_aval_array
 
 # TODO(jakevdp) replace this with true inheritance at the C++ level.
 basearray.Array.register(ArrayImpl)

jax/_src/core.py

@@ -1388,7 +1388,7 @@ def lattice_join(x, y):
 
 def valid_jaxtype(x) -> bool:
   try:
-    abstractify(x)
+    concrete_aval(x)
   except TypeError:
     return False
   else:
@@ -1400,9 +1400,35 @@ def check_valid_jaxtype(x):
       f"Value {x!r} of type {type(x)} is not a valid JAX type")
 
 
-def abstractify(x):
+# TODO(jakevdp): merge concrete_aval and abstractify to the extent possible.
+# This is tricky because concrete_aval includes sharding information, and
+# abstractify does not; further, because abstractify is in the dispatch path,
+# performance is important and simply adding sharding there is not an option.
+def concrete_aval(x):
+  # This differs from abstractify below in that the abstract values
+  # include sharding where applicable. Historically (before stackless)
+  # the returned avals were concrete, but after the stackless change
+  # this returns ShapedArray like abstractify.
+  # Rules are registered in pytype_aval_mappings.
   for typ in type(x).__mro__:
-    aval_fn = pytype_aval_mappings.get(typ)
+    handler = pytype_aval_mappings.get(typ)
+    if handler: return handler(x)
+  if hasattr(x, '__jax_array__'):
+    return concrete_aval(x.__jax_array__())
+  raise TypeError(f"Value {x!r} with type {type(x)} is not a valid JAX "
+                   "type")
+
+
+def abstractify(x):
+  # Historically, this was called xla.abstractify. It differs from
+  # concrete_aval in that it excludes sharding information, and
+  # uses a more performant path for accessing avals. Rules are
+  # registered in xla_pytype_aval_mappings.
+  typ = type(x)
+  aval_fn = xla_pytype_aval_mappings.get(typ)
+  if aval_fn: return aval_fn(x)
+  for typ in typ.__mro__:
+    aval_fn = xla_pytype_aval_mappings.get(typ)
     if aval_fn: return aval_fn(x)
   if hasattr(x, '__jax_array__'):
     return abstractify(x.__jax_array__())
@@ -1413,7 +1439,7 @@ def get_aval(x):
   if isinstance(x, Tracer):
     return x.aval
   else:
-    return abstractify(x)
+    return concrete_aval(x)
 
 get_type = get_aval
 
@@ -1809,6 +1835,7 @@ def to_tangent_aval(self):
                         self.weak_type)
 
 pytype_aval_mappings: dict[type, Callable[[Any], AbstractValue]] = {}
+xla_pytype_aval_mappings: dict[type, Callable[[Any], AbstractValue]] = {}
 
 
 class DArray:
@@ -1865,6 +1892,7 @@ def data(self):
 
 pytype_aval_mappings[DArray] = \
     lambda x: DShapedArray(x._aval.shape, x._aval.dtype, x._aval.weak_type)
+xla_pytype_aval_mappings[DArray] = lambda x: x._aval
 
 @dataclass(frozen=True)
 class bint(dtypes.ExtendedDType):
@@ -1897,6 +1925,7 @@ def __getitem__(self, idx): return get_aval(self)._getitem(self, idx)
   def __setitem__(self, idx, x): return get_aval(self)._setitem(self, idx, x)
   def __repr__(self) -> str: return 'Mutable' + repr(self[...])
 pytype_aval_mappings[MutableArray] = lambda x: x._aval
+xla_pytype_aval_mappings[MutableArray] = lambda x: x._aval
 
 def mutable_array(init_val):
   return mutable_array_p.bind(init_val)
@@ -1952,6 +1981,7 @@ def __init__(self, buf):
   def block_until_ready(self):
     self._buf.block_until_ready()
 pytype_aval_mappings[Token] = lambda _: abstract_token
+xla_pytype_aval_mappings[Token] = lambda _: abstract_token
 
 
 # TODO(dougalm): Deprecate these. They're just here for backwards compat.

jax/_src/export/shape_poly.py

@@ -1205,6 +1205,7 @@ def _geq_decision(e1: DimSize, e2: DimSize, cmp_str: Callable[[], str]) -> bool:
       f"Symbolic dimension comparison {cmp_str()} is inconclusive.{describe_scope}")
 
 core.pytype_aval_mappings[_DimExpr] = _DimExpr._get_aval
+core.xla_pytype_aval_mappings[_DimExpr] = _DimExpr._get_aval
 dtypes._weak_types.append(_DimExpr)
 
 def _convertible_to_int(p: DimSize) -> bool:

jax/_src/interpreters/xla.py

@@ -146,6 +146,13 @@ def _canonicalize_python_scalar_dtype(typ, x):
 canonicalize_dtype_handlers[core.DArray] = identity
 canonicalize_dtype_handlers[core.MutableArray] = identity
 
+# TODO(jakevdp): deprecate and remove this.
+def abstractify(x) -> Any:
+  return core.abstractify(x)
+
+# TODO(jakevdp): deprecate and remove this.
+pytype_aval_mappings: dict[Any, Callable[[Any], core.AbstractValue]] = core.xla_pytype_aval_mappings
+
 initial_style_primitives: set[core.Primitive] = set()
 
 def register_initial_style_primitive(prim: core.Primitive):

jax/_src/prng.py

@@ -463,6 +463,7 @@ def __hash__(self) -> int:
 
 
 core.pytype_aval_mappings[PRNGKeyArray] = lambda x: x.aval
+core.xla_pytype_aval_mappings[PRNGKeyArray] = lambda x: x.aval
 
 xla.canonicalize_dtype_handlers[PRNGKeyArray] = lambda x: x
 

jax/core.py

@@ -122,7 +122,7 @@
                _src_core.call_p),
     "closed_call_p": ("jax.core.closed_call_p is deprecated. Use jax.extend.core.primitives.closed_call_p",
                       _src_core.closed_call_p),
-    "concrete_aval": ("jax.core.concrete_aval is deprecated.", _src_core.abstractify),
+    "concrete_aval": ("jax.core.concrete_aval is deprecated.", _src_core.concrete_aval),
     "dedup_referents": ("jax.core.dedup_referents is deprecated.", _src_core.dedup_referents),
     "escaped_tracer_error": ("jax.core.escaped_tracer_error is deprecated.",
                              _src_core.escaped_tracer_error),
@@ -207,7 +207,7 @@
   axis_frame = _src_core.axis_frame
   call_p = _src_core.call_p
   closed_call_p = _src_core.closed_call_p
-  concrete_aval = _src_core.abstractify
+  concrete_aval = _src_core.concrete_aval
   dedup_referents = _src_core.dedup_referents
   escaped_tracer_error = _src_core.escaped_tracer_error
   extend_axis_env_nd = _src_core.extend_axis_env_nd

jax/interpreters/xla.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 from jax._src.interpreters.xla import (
+  abstractify as abstractify,
   canonicalize_dtype as canonicalize_dtype,
   canonicalize_dtype_handlers as canonicalize_dtype_handlers,
+  pytype_aval_mappings as pytype_aval_mappings,
 )
 
 from jax._src.dispatch import (
@@ -25,19 +27,8 @@
 Backend = _xc._xla.Client
 del _xc
 
-from jax._src import core as _src_core
-
 # Deprecations
 _deprecations = {
-    # Added 2024-12-17
-    "abstractify": (
-        "jax.interpreters.xla.abstractify is deprecated.",
-        _src_core.abstractify
-    ),
-    "pytype_aval_mappings": (
-        "jax.interpreters.xla.pytype_aval_mappings is deprecated.",
-        _src_core.pytype_aval_mappings
-    ),
     # Finalized 2024-10-24; remove after 2025-01-24
     "xb": (
         ("jax.interpreters.xla.xb was removed in JAX v0.4.36. "
@@ -53,13 +44,6 @@
     ),
 }
 
-import typing as _typing
 from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
-if _typing.TYPE_CHECKING:
-  abstractify = _src_core.abstractify
-  pytype_aval_mappings = _src_core.pytype_aval_mappings
-else:
-  __getattr__ = _deprecation_getattr(__name__, _deprecations)
+__getattr__ = _deprecation_getattr(__name__, _deprecations)
 del _deprecation_getattr
-del _typing
-del _src_core

tests/lax_test.py

@@ -3959,6 +3959,8 @@ def setUp(self):
     core.pytype_aval_mappings[FooArray] = \
         lambda x: core.ShapedArray(x.shape, FooTy())
     xla.canonicalize_dtype_handlers[FooArray] = lambda x: x
+    core.xla_pytype_aval_mappings[FooArray] = \
+        lambda x: core.ShapedArray(x.shape, FooTy())
     pxla.shard_arg_handlers[FooArray] = shard_foo_array_handler
     mlir._constant_handlers[FooArray] = foo_array_constant_handler
     mlir.register_lowering(make_p, mlir.lower_fun(make_lowering, False))
@@ -3971,6 +3973,7 @@ def setUp(self):
   def tearDown(self):
     del core.pytype_aval_mappings[FooArray]
     del xla.canonicalize_dtype_handlers[FooArray]
+    del core.xla_pytype_aval_mappings[FooArray]
     del mlir._constant_handlers[FooArray]
     del mlir._lowerings[make_p]
     del mlir._lowerings[bake_p]