Removed jax._src.raise_to_shaped

It is just an identity after the "stackless" rewrite.

PiperOrigin-RevId: 745042532

Author: superbobry

jax/_src/core.py

@@ -2229,11 +2229,6 @@ def block_until_ready(self):
 pytype_aval_mappings[Token] = lambda _: abstract_token
 
 
-# TODO(dougalm): Deprecate these. They're just here for backwards compat.
-def raise_to_shaped(aval):
-  return aval
-raise_to_shaped_mappings: dict[type, Callable] = {}
-
 ### Operations on shapes and dimension sizes.
 
 class InconclusiveDimensionOperation(Exception):

jax/_src/pallas/fuser/fusable.py

@@ -29,10 +29,6 @@
 fusable_p.multiple_results = True
 
 
-def _get_aval(x):
-  return jax_core.raise_to_shaped(jax_core.get_aval(x))
-
-
 def _make_trivial_fusion(x: jax.Array) -> fusion_lib.Fusion:
   return fusion_lib.Fusion(
       func=lambda: x,
@@ -53,7 +49,7 @@ def wrapped(*args):
       flat_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
           lu.wrap_init(wrapped, debug_info=debug_info), in_tree
       )
-      flat_avals = [_get_aval(x) for x in flat_args]
+      flat_avals = [jax_core.get_aval(x) for x in flat_args]
       jaxpr, _, consts, _ = pe.trace_to_jaxpr_dynamic(flat_fun, flat_avals)
       out_tree = out_tree_thunk()
       out = fusable_p.bind(

jax/_src/pallas/fuser/jaxpr_fusion.py

@@ -28,10 +28,6 @@
 from jax._src.pallas.fuser.fusable import fusable_p
 
 
-def _get_aval(x):
-  return jax_core.raise_to_shaped(jax_core.get_aval(x))
-
-
 def fuse(f=None, *, physicalize: bool = False, debug: bool = False):
   """Fuses a function into a single fusable.
 
@@ -52,7 +48,7 @@ def wrapper(*args, **kwargs):
       flat_fun, out_tree_thunk = api_util.flatten_fun(
           lu.wrap_init(f, debug_info=debug_info), in_tree
       )
-      flat_avals = [_get_aval(x) for x in flat_args]
+      flat_avals = [jax_core.get_aval(x) for x in flat_args]
       jaxpr, _, consts, _ = pe.trace_to_jaxpr_dynamic(flat_fun, flat_avals)
       if debug:
         print("Jaxpr before fusion:")

tests/api_test.py

@@ -5024,7 +5024,7 @@ def g(x):
 
     # Make sure that introducing constants in vmap works.
     constant_introducing_p = core.Primitive('introduce_constant')
-    constant_introducing_p.def_abstract_eval(core.raise_to_shaped)
+    constant_introducing_p.def_abstract_eval(lambda x: x)
     def _constant_introducing_batcher(xs, ds):
       (x,), (d,) = xs, ds
       return (x + np.arange(x.size, dtype=x.dtype).reshape(x.shape)), d

tests/state_test.py

@@ -792,7 +792,7 @@ def body(i, st):
       lax.fori_loop(0, 5, body, init_val=())
       return a_ref[...], b_ref[...]
 
-    ref = lambda x: AbstractRef(core.raise_to_shaped(core.get_aval(x)))
+    ref = lambda x: AbstractRef(core.get_aval(x))
     f_jaxpr = jax.make_jaxpr(f)(ref(1.), ref(2.))
     jaxpr, _ = discharge_state(f_jaxpr.jaxpr, (), should_discharge=[False, True])
     # Effects on y_ref were discharged away but not the effects on x_ref
@@ -1139,7 +1139,7 @@ def false_fun():
         y_ref[...] = 2.
       lax.cond(pred, true_fun, false_fun)
       return x_ref[...], y_ref[...]
-    ref = lambda x: AbstractRef(core.raise_to_shaped(core.get_aval(x)))
+    ref = lambda x: AbstractRef(core.get_aval(x))
     f_jaxpr = jax.make_jaxpr(f0)(False, ref(3.), ref(4.))
     jaxpr, _ = discharge_state(f_jaxpr.jaxpr, (), should_discharge=[False, False, True])
     # Effects on y_ref were discharged away but not the effects on x_ref