Make MetaTy work with HiJAX. Before this HiTypes (with values) were being held in Traced without converting it to MetaTys.

This change makes MetaTys work with HiTypes.

Co-authored-by: Matthew Johnson <mattjj@google.com>
PiperOrigin-RevId: 826153440

Author: 2 people

jax/_src/pjit.py

@@ -16,7 +16,7 @@
 
 from collections import defaultdict
 from collections.abc import Callable, Sequence, Iterable
-import dataclasses
+from dataclasses import dataclass, replace
 from functools import partial
 import inspect
 import logging
@@ -153,7 +153,7 @@ def _python_pjit_helper(fun: Callable, jit_info: PjitInfo, *args, **kwargs):
   except stages.DeviceAssignmentMismatchError as e:
     fails, = e.args
     fun_name = getattr(fun, '__qualname__', getattr(fun, '__name__', str(fun)))
-    arg_types = map(partial(convert_to_lower_type, False), args_flat)
+    arg_types = map(convert_to_metaty, args_flat)
     msg = stages._device_assignment_mismatch_error(
         fun_name, fails, arg_types, 'jit', p.arg_names)
     raise ValueError(msg) from None
@@ -304,8 +304,7 @@ def cache_miss(*args, **kwargs):
 @api_boundary
 def jit_trace(jit_func, *args, **kwargs) -> stages.Traced:
   p, args_flat = _infer_params(jit_func._fun, jit_func._jit_info, args, kwargs)
-  arg_types = map(partial(convert_to_lower_type, p.params['jaxpr'].is_high),
-                  args_flat)
+  arg_types = map(convert_to_metaty, args_flat)
   return stages.Traced(arg_types, p.params, p.in_tree, p.out_tree, p.consts)
 
 @api_boundary
@@ -1253,7 +1252,7 @@ def _qdd_cache_update(fun, in_type, i, consts, aval_qdds):
                   if aval_qdd.has_qdd])
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclass(frozen=True)
 class IgnoreKey:
   val: Any
   def __hash__(self):
@@ -1488,7 +1487,7 @@ def _resolve_in_shardings(args, pjit_in_shardings: Sequence[PjitSharding]
     if isinstance(pjit_in_s, UnspecifiedValue):
       resolved_in_shardings.append(finalize_arg_sharding(arg_s, committed))
     else:
-      if (arg.is_nd_array and not pjit_in_s.is_fully_replicated and  # type: ignore[union-attr]
+      if (arg.is_np_array and not pjit_in_s.is_fully_replicated and  # type: ignore[union-attr]
           xb.process_count() > 1):
         raise ValueError(
             'Passing non-trivial shardings for numpy '
@@ -1542,13 +1541,15 @@ def _resolve_and_lower(
 _pgle_profiler_dict = weakref.WeakKeyDictionary()  # type: ignore
 
 
-@dataclasses.dataclass(frozen=True)
-class LowerType:
+@dataclass(frozen=True)
+class MetaTy:
+  aval: Any
   sharding: Any
   format: Any
   committed: bool
-  aval: Any
-  is_nd_array: bool
+  is_np_array: bool
+
+  replace = replace  # type: ignore
 
   @property
   def shape(self):
@@ -1559,15 +1560,13 @@ def ndim(self):
     return self.aval.ndim
 
 
-def convert_to_lower_type(is_high, arg):
-  if is_high:
-    return arg
+def convert_to_metaty(arg):
   arg_sharding = getattr(arg, 'sharding', None)
   arg_format = getattr(arg, 'format', None)
   arg_committed = getattr(arg, '_committed', True)
   aval = core.shaped_abstractify(arg)
-  is_nd_array = isinstance(arg, np.ndarray)
-  return LowerType(arg_sharding, arg_format, arg_committed, aval, is_nd_array)
+  is_np_array = isinstance(arg, np.ndarray)
+  return MetaTy(aval, arg_sharding, arg_format, arg_committed, is_np_array)
 
 
 def _pjit_call_impl_python(
@@ -1597,7 +1596,7 @@ def _pjit_call_impl_python(
   compiler_options_kvs = compiler_options_kvs + tuple(pgle_compile_options.items())
   # Passing mutable PGLE profile here since it should be extracted by JAXPR to
   # initialize the fdo_profile compile option.
-  arg_types = map(partial(convert_to_lower_type, jaxpr.is_high), args)
+  arg_types = map(convert_to_metaty, args)
   computation = _resolve_and_lower(
       arg_types, jaxpr=jaxpr, in_shardings=in_shardings,
       out_shardings=out_shardings, in_layouts=in_layouts,

jax/_src/stages.py

@@ -34,7 +34,6 @@
 import enum
 from collections.abc import Sequence
 from dataclasses import dataclass
-from functools import partial
 import itertools as it
 from typing import Any, NamedTuple, Protocol, Union, runtime_checkable
 
@@ -404,10 +403,10 @@ class Traced(Stage):
   traced representation with the remaining information needed to later
   lower, compile, and execute it.
   """
-  __slots__ = ['_args_flat', '_params', '_in_tree', 'out_tree', '_consts']
+  __slots__ = ['_meta_tys_flat', '_params', '_in_tree', 'out_tree', '_consts']
 
-  def __init__(self, args_flat, params, in_tree, out_tree, consts):
-    self._args_flat = args_flat
+  def __init__(self, meta_tys_flat, params, in_tree, out_tree, consts):
+    self._meta_tys_flat = meta_tys_flat
     self._params = params
     self._in_tree = in_tree
     self.out_tree = out_tree
@@ -425,13 +424,12 @@ def out_avals(self):
 
   def fall(self):
     if not self.jaxpr.is_high:
-      return Fallen(self._args_flat, self._params, self._in_tree, self.out_tree,
-                    (self._in_tree, self.jaxpr.in_avals),
+      return Fallen(self._meta_tys_flat, self._params, self._in_tree,
+                    self.out_tree, (self._in_tree, self.jaxpr.in_avals),
                     (self.out_tree, self.jaxpr.out_avals),
                     self._consts)
 
     # TODO(mattjj): when pmap is deleted, merge with pjit.py BUILD rule
-    from jax._src.pjit import convert_to_lower_type  # type: ignore
     from jax._src.interpreters import partial_eval as pe  # type:ignore
     hi_jaxpr = self.jaxpr
     _, closed_over_himutables = pe.convert_const_himutables(hi_jaxpr)
@@ -440,11 +438,11 @@ def fall(self):
     in_tree = lojax_pytree(hi_jaxpr.in_aval_qdds, self._in_tree)
     out_tree = lojax_pytree(hi_jaxpr.out_avals, self.out_tree)
     params = dict(lojax_expand_params(hi_jaxpr, self._params), jaxpr=lo_jaxpr)
-    lo_args = [lo_val for aval, x in zip(hi_jaxpr.in_aval_qdds, self._args_flat)
-               for lo_val in (aval.read_loval(x) if aval.has_qdd
-                              else aval.lower_val(x))]
-    lo_arg_types = map(partial(convert_to_lower_type, False), lo_args)
-    return Fallen(lo_arg_types, params, in_tree, out_tree,
+    lo_meta_tys = [mty.replace(aval=lo_ty)
+                   for mty, aq in zip(self._meta_tys_flat, hi_jaxpr.in_aval_qdds)
+                   for lo_ty in (mty.aval.lo_ty_qdd(aq.qdd)
+                                 if mty.aval.has_qdd else mty.aval.lo_ty())]
+    return Fallen(lo_meta_tys, params, in_tree, out_tree,
                   (self._in_tree, hi_jaxpr.final_aval_qdds),
                   (self.out_tree, hi_jaxpr.out_avals),
                   self._consts)
@@ -472,12 +470,12 @@ def lojax_pytree(hi_avals, tree):
 
 class Fallen(Stage):
   """True leader of the Decepticons."""
-  __slots__ = ['_args_flat', '_params', '_in_tree', 'out_tree',
+  __slots__ = ['_meta_tys_flat', '_params', '_in_tree', 'out_tree',
                '_consts', '_in_types', '_out_types']
 
-  def __init__(self, args_flat, params, in_tree, out_tree, in_types, out_types,
+  def __init__(self, meta_tys_flat, params, in_tree, out_tree, in_types, out_types,
                consts):
-    self._args_flat = args_flat
+    self._meta_tys_flat = meta_tys_flat
     self._params = params
     self._in_tree = in_tree
     self.out_tree = out_tree
@@ -503,12 +501,12 @@ def lower(self, *, lowering_platforms: tuple[str, ...] | None = None,
     try:
       from jax._src.pjit import _resolve_and_lower  # type: ignore
       lowering = _resolve_and_lower(
-          self._args_flat, **self._params, lowering_platforms=lowering_platforms,
+          self._meta_tys_flat, **self._params, lowering_platforms=lowering_platforms,
           lowering_parameters=_private_parameters, pgle_profiler=None)
     except DeviceAssignmentMismatchError as e:
       fails, = e.args
       msg = _device_assignment_mismatch_error(
-          self._params['name'], fails, self._args_flat, 'jit',
+          self._params['name'], fails, self._meta_tys_flat, 'jit',
           self.jaxpr.debug_info.safe_arg_names(len(self.jaxpr.in_avals)))
       raise ValueError(msg) from None
     return Lowered(lowering, self.args_info, self.out_tree,