Remove eval_shape as a method on Traced and rather use .out_info because .trace already evals.

Forward `jit(f).eval_shape(*args)` to `jit(f).trace(*args).out_info`

Co-authored-by: Roy Frostig <[email protected]>
PiperOrigin-RevId: 786896497

Author: 2 people

jax/_src/api.py

@@ -2842,10 +2842,10 @@ def eval_shape(fun, *args, **kwargs):
   float32
   """
   if type(fun) is xc._xla.PjitFunction:
-    return fun.trace(*args, **kwargs).eval_shape()  # type: ignore
+    return fun.trace(*args, **kwargs).out_info  # type: ignore
   try: hash(fun)
   except TypeError: fun = partial(fun)
-  return jit(fun).trace(*args, **kwargs).eval_shape()
+  return jit(fun).trace(*args, **kwargs).out_info
 
 
 def named_call(

jax/_src/pjit.py

@@ -313,7 +313,7 @@ def jit_lower(jit_func, *args, **kwargs):
 
 @api_boundary
 def jit_eval_shape(jit_func, *args, **kwargs):
-  return jit_trace(jit_func, *args, **kwargs).eval_shape()
+  return jit_trace(jit_func, *args, **kwargs).out_info
 
 def jit_evict_fn(self):
   self._clear_cache()

jax/_src/stages.py

@@ -709,7 +709,18 @@ def __init__(self, jaxpr: core.ClosedJaxpr, args_info, fun_name, out_tree,
 
   @property
   def out_info(self):
-    return self.eval_shape()
+    out_shardings = [None if isinstance(s, UnspecifiedValue) else s
+                     for s in self._params_out_shardings]
+    out = []
+    for a, out_s in zip(self.jaxpr.out_avals, out_shardings):
+      s = (a.sharding if a.sharding.mesh._are_all_axes_explicit else out_s
+           if out_s is None else out_s)
+      # TODO(yashkatariya): Add `Layout` to SDS.
+      out.append(
+          core.ShapeDtypeStruct(
+              a.shape, a.dtype, sharding=s, weak_type=a.weak_type,
+              vma=(a.vma if config._check_vma.value else None)))
+    return tree_util.tree_unflatten(self._out_tree, out)
 
   def lower(self, *, lowering_platforms: tuple[str, ...] | None = None,
             _private_parameters: mlir.LoweringParameters | None = None):
@@ -727,20 +738,6 @@ def lower(self, *, lowering_platforms: tuple[str, ...] | None = None,
       raise ValueError(msg) from None
     return Lowered(lowering, self.args_info, self._out_tree)
 
-  def eval_shape(self):
-    out_shardings = [None if isinstance(s, UnspecifiedValue) else s
-                     for s in self._params_out_shardings]
-    out = []
-    for a, out_s in zip(self.jaxpr.out_avals, out_shardings):
-      s = (a.sharding if a.sharding.mesh._are_all_axes_explicit else out_s
-           if out_s is None else out_s)
-      # TODO(yashkatariya): Add `Layout` to SDS.
-      out.append(
-          core.ShapeDtypeStruct(
-              a.shape, a.dtype, sharding=s, weak_type=a.weak_type,
-              vma=(a.vma if config._check_vma.value else None)))
-    return tree_util.tree_unflatten(self._out_tree, out)
-
 
 @runtime_checkable
 class Wrapped(Protocol):