Try to add constants directly into core.Jaxpr

DO_NOT_SUBMIT

Author: gnecula

jax/_src/api.py

@@ -2425,7 +2425,7 @@ def make_jaxpr_f(*args, **kwargs):
     if traced._num_consts:
       consts, _ = split_list(traced._args_flat, [traced._num_consts])
       jaxpr_ = pe.convert_invars_to_constvars(traced.jaxpr.jaxpr,
-                                              traced._num_consts)
+                                              traced._num_consts, consts)
       jaxpr = core.ClosedJaxpr(jaxpr_, consts)
     else:
       jaxpr = traced.jaxpr

jax/_src/core.py

@@ -88,7 +88,7 @@
 
 class Jaxpr:
   __slots__ = ['__weakref__', '_constvars', '_invars', '_outvars', '_eqns',
-               '_effects', '_debug_info', '_is_high']
+               '_effects', '_debug_info', '_is_high', '_consts']
 
   _constvars: list[Var]
   _invars: list[Var]
@@ -126,14 +126,19 @@ def debug_info(self) -> DebugInfo:
   def is_high(self) -> bool:
     return self._is_high
 
+  @property
+  def consts(self) -> list[Any]:
+    return self._consts
+
   def __init__(self, constvars: Sequence[Var], invars: Sequence[Var],
                outvars: Sequence[Atom], eqns: Sequence[JaxprEqn],
                effects: Effects = no_effects,
                # We want all calls to pass a DebugInfo object, but for backwards
                # compatibility we have to allow calls when the debug_info
                # is missing.
                debug_info: DebugInfo = None,  # type: ignore[annotation-type-mismatch,assignment]
-               is_high: bool = False,
+               is_high: bool = False, *,
+               consts: Sequence[Any] = (),
                ):
     """
     Args:
@@ -146,6 +151,7 @@ def __init__(self, constvars: Sequence[Var], invars: Sequence[Var],
       effects: set of effects. The effects on a jaxpr are a superset of the
         union of the effects for each equation.
       debug_info: debugging information.
+      consts: the constant values corresponding to the constvars
     """
     self._constvars = list(constvars)
     self._invars = list(invars)
@@ -159,7 +165,10 @@ def __init__(self, constvars: Sequence[Var], invars: Sequence[Var],
     # assert (len(debug_info.arg_names) == len(invars)), (debug_info, invars)
     # assert (len(debug_info.result_paths) == len(outvars)), (debug_info, outvars)
     self._is_high = is_high
-    num_vars = len(constvars) + len(invars)
+    self._consts = list(consts)
+    if len(constvars) != len(consts):
+      assert False, (constvars, consts)  # DO_NOT_SUBMIT
+    assert len(constvars) == len(consts), (constvars, consts)
 
   def __str__(self):
     return str(self.pretty_print())
@@ -187,6 +196,7 @@ def replace(self, **kwargs):
         effects=kwargs.pop("effects", self.effects),
         debug_info=kwargs.pop("debug_info", self.debug_info),
         is_high=kwargs.pop("is_high", self.is_high),
+        consts=kwargs.pop("consts", self.consts),
     )
     if kwargs:
       raise ValueError(f"Unknown keyword arguments: {kwargs}")
@@ -223,11 +233,14 @@ class ClosedJaxpr:
   jaxpr = property(lambda self: self._jaxpr)
   consts = property(lambda self: self._consts)
 
-  def __init__(self, jaxpr: Jaxpr, consts: Sequence):
+  def __init__(self, jaxpr: Jaxpr, consts: Sequence[Value]):
     assert len(consts) == len(jaxpr.constvars)
     # assert not any(isinstance(c, Tracer) for c in consts)  # TODO(mattjj): enable
     self._jaxpr = jaxpr
     self._consts = list(consts)
+    if jaxpr.consts != self._consts:  # DO_NOT_SUBMIT
+      assert False, (jaxpr.consts, self._consts)  # TODO(necula): remove, when we remove ClosedJaxpr
+    assert jaxpr.consts == self._consts, (jaxpr.consts, self._consts)  # TODO(necula): remove, when we remove ClosedJaxpr
 
   @property
   def in_avals(self):

jax/_src/interpreters/ad.py

@@ -1224,7 +1224,7 @@ def rearrange_binders(jaxpr: core.ClosedJaxpr, primals_in, tangents_in, primals_
       (*constvars, *new_invars), (*constvars, *jaxpr.jaxpr.invars),
       jaxpr.jaxpr.effects)
   new_jaxpr = core.Jaxpr(constvars, new_invars, new_outvars, jaxpr.jaxpr.eqns,
-                         new_effects, new_debug_info)
+                         new_effects, new_debug_info, consts=jaxpr.consts)
   return core.ClosedJaxpr(new_jaxpr, jaxpr.consts)
 
 def _perm(primal_counts: Sequence[int], tangent_counts: Sequence[int],

jax/_src/interpreters/partial_eval.py

@@ -880,7 +880,7 @@ def sort_key(t):
   jaxpr_effects = make_jaxpr_effects(const_vars, invars, outvars, eqns)
   jaxpr = Jaxpr(const_vars, invars,  # type: ignore[arg-type]
                 outvars, eqns, jaxpr_effects,
-                debug_info)
+                debug_info, consts=const_vals)
   config.enable_checks.value and core.check_jaxpr(jaxpr)
   # del getvar  # needed to avoid cyclic-reference closure, apparently!
   return jaxpr, const_vals, env_vals
@@ -897,35 +897,27 @@ def convert_constvars_jaxpr(jaxpr: Jaxpr) -> Jaxpr:
   dbg = jaxpr.debug_info._replace(
       arg_names=("",) * len(jaxpr.constvars) + jaxpr.debug_info.arg_names)
   lifted_jaxpr = jaxpr.replace(
-      constvars=(), invars=jaxpr.constvars + jaxpr.invars, debug_info=dbg)
+      constvars=(), invars=jaxpr.constvars + jaxpr.invars, debug_info=dbg, consts=())
   config.enable_checks.value and core.check_jaxpr(lifted_jaxpr)
   return lifted_jaxpr
 
-@weakref_lru_cache
-def convert_invars_to_constvars(jaxpr: Jaxpr, n: int) -> Jaxpr:
-  """Move n invars to constvars. Like an inverse of convert_constvars_jaxpr."""
+# DO_NOT_SUBMIT
+# @weakref_lru_cache
+def convert_invars_to_constvars(jaxpr: Jaxpr, n: int, consts: Sequence[core.Value]) -> Jaxpr:
+  """Move the first n invars to constvars. Like an inverse of convert_constvars_jaxpr."""
   if n == 0:
     return jaxpr.replace()  # 'return jaxpr' would create cache reference cycle
   config.enable_checks.value and core.check_jaxpr(jaxpr)
+  assert n == len(consts)
+  assert 0 == len(jaxpr.constvars)
   constvars, invars = split_list(jaxpr.invars, [n])
   dbg = jaxpr.debug_info._replace(
       arg_names=jaxpr.debug_info.arg_names[n:])
   lifted_jaxpr = jaxpr.replace(constvars=tuple(constvars), invars=invars,
-                               debug_info=dbg)
+                               debug_info=dbg, consts=consts)
   config.enable_checks.value and core.check_jaxpr(lifted_jaxpr)
   return lifted_jaxpr
 
-def convert_envvars_to_constvars(jaxpr: Jaxpr, num_env_vars: int) -> Jaxpr:
-  if any(isinstance(eff, effects.JaxprInputEffect) for eff in jaxpr.effects):
-    raise NotImplementedError
-  config.enable_checks.value and core.check_jaxpr(jaxpr)
-  env_vars, invars = split_list(jaxpr.invars, [num_env_vars])
-  converted_jaxpr = jaxpr.replace(constvars=jaxpr.constvars + env_vars,
-                                  invars=invars)
-  config.enable_checks.value and core.check_jaxpr(converted_jaxpr)
-  return converted_jaxpr
-
-
 def partial_eval_jaxpr_nounits(
     jaxpr: ClosedJaxpr, unknowns: Sequence[bool],
     instantiate: bool | Sequence[bool],
@@ -1459,7 +1451,9 @@ def dce_jaxpr_consts(jaxpr: Jaxpr, used_outputs: Sequence[bool],
   new_jaxpr, used_inputs_ = dce_jaxpr(jaxpr_, used_outputs, instantiate)
   used_consts, used_inputs = split_list(used_inputs_, [len(jaxpr.constvars)])
   if sum(used_consts):
-    new_jaxpr = convert_invars_to_constvars(new_jaxpr, sum(used_consts))
+    new_consts = tuple(c for u, c in zip(used_consts, jaxpr.consts) if u)
+    new_jaxpr = convert_invars_to_constvars(new_jaxpr, sum(used_consts),
+                                            new_consts)
   return new_jaxpr, used_consts, used_inputs
 
 
@@ -1805,13 +1799,14 @@ def to_jaxpr(
       v.final_qdd = qdd.cur_val
 
     jaxpr = Jaxpr(constvars, invars, outvars, self.eqns, jaxpr_effects,
-                  debug_info, self.is_high)
+                  debug_info, self.is_high, consts=constvals)
     jaxpr, constvals = _drop_unused_vars(jaxpr, constvals)
     init_trees = [tree_structure(init_val) for init_val in self.attrs_inits]
     return jaxpr, list(constvals), zip(init_trees, end_trees, self.attrs_tracked)
 
   def to_jaxpr2(self, out_tracers: Sequence[core.Tracer],
-                debug_info: core.DebugInfo):
+                debug_info: core.DebugInfo
+                ) -> tuple[core.Jaxpr, tuple[core.AbstractValue, ...]]:
     # It's not necessary, but we keep the tracer-to-var mapping injective:
     vars = [v for v in self.tracer_to_var.values() if not isinstance(v, Literal)]
     assert len(vars) == len(set(vars))
@@ -1820,7 +1815,7 @@ def to_jaxpr2(self, out_tracers: Sequence[core.Tracer],
     jaxpr_effects = make_jaxpr_effects(constvars, self.invars, expl_outvars,
                                         self.eqns)
     jaxpr = Jaxpr(constvars, self.invars, expl_outvars, self.eqns,
-                  jaxpr_effects, debug_info)
+                  jaxpr_effects, debug_info, consts=constvals)
     # We can't run check_jaxpr until after we normalize.
     jaxpr, constvals = _drop_unused_vars(jaxpr, constvals)
     jaxpr, out_type = _add_implicit_outputs(jaxpr)
@@ -1886,9 +1881,10 @@ def vars(atom: Atom) -> list[Var]:
     used.update(v for atom in eqn.invars for v in vars(atom))
   cvars, constvals = unzip2(
       (v, val) for v, val in zip(jaxpr.constvars, constvals) if v in used)
-  jaxpr._constvars = list(cvars)
-  jaxpr._effects = make_jaxpr_effects(jaxpr.constvars, jaxpr.invars,
-                                      jaxpr.outvars, jaxpr.eqns)
+  jaxpr = jaxpr.replace(constvars=list(cvars),
+                        effects=make_jaxpr_effects(jaxpr.constvars, jaxpr.invars,
+                                      jaxpr.outvars, jaxpr.eqns),
+                        consts=list(constvals))
   return jaxpr, list(constvals)
 
 
@@ -2831,5 +2827,5 @@ def convert_const_himutables(jaxpr):
   effects = make_jaxpr_effects(constvars, invars, jaxpr.jaxpr.outvars,
                                jaxpr.jaxpr.eqns)
   new_jaxpr = jaxpr.jaxpr.replace(constvars=constvars, invars=invars,
-                                  effects=effects)
+                                  effects=effects, consts=constvals)
   return jaxpr.replace(jaxpr=new_jaxpr, consts=constvals), in_mutables

jax/_src/interpreters/pxla.py

@@ -1734,6 +1734,8 @@ def _dce_jaxpr(closed_jaxpr, api_name, fun_name,
   in_avals = closed_jaxpr.in_avals
 
   if (keep_unused or auto_spmd_lowering or
+      # Don't drop inputs with symbolic shapes, because we may need them
+      # to infer the symbolic dimensions at call sites.
       any(hasattr(a, "shape") and not core.is_constant_shape(a.shape)
           for a in in_avals)):
     kept_var_idx = set(range(len(in_avals)))

jax/_src/lax/control_flow/common.py

@@ -175,13 +175,18 @@ def _initial_style_jaxprs_with_common_consts(
     canonical_non_ref_indices.append(tuple(non_ref_indices))
 
   consts = [*canonical_refs, *canonical_non_refs]
-  jaxprs = tuple(_pad_jaxpr_constvars(jaxpr, i, (*canonical_ref_avals,), (*canonical_ref_indices,), (*canonical_non_ref_avals,), (*canonical_non_ref_indices,))
+  jaxprs = tuple(_pad_jaxpr_constvars(jaxpr, i, consts,
+                                      (*canonical_ref_avals,), (*canonical_ref_indices,), (*canonical_non_ref_avals,), (*canonical_non_ref_indices,))
                  for i, jaxpr in enumerate(jaxprs))
   return jaxprs, consts, all_out_trees
 
-@weakref_lru_cache
-def _pad_jaxpr_constvars(jaxpr, i, canonical_ref_avals, canonical_ref_indices,
-                         canonical_non_ref_avals, canonical_non_ref_indices):
+# DO_NOT_SUBMIT
+# @weakref_lru_cache
+def _pad_jaxpr_constvars(jaxpr: core.Jaxpr, i,
+                         consts: Sequence[core.Value],
+                         canonical_ref_avals, canonical_ref_indices,
+                         canonical_non_ref_avals, canonical_non_ref_indices
+                         ) -> core.ClosedJaxpr:
   is_ref = [isinstance(v.aval, state.AbstractRef) for v in jaxpr.constvars]
   nonref_constvars, ref_constvars = partition_list(is_ref, jaxpr.constvars)
   padded_ref_constvars  = map(core.Var, canonical_ref_avals)
@@ -191,7 +196,7 @@ def _pad_jaxpr_constvars(jaxpr, i, canonical_ref_avals, canonical_ref_indices,
   for canonical_id, non_ref_var in zip(canonical_non_ref_indices[i], nonref_constvars):
     padded_non_ref_constvars[canonical_id] = non_ref_var
   constvars = [*padded_ref_constvars, *padded_non_ref_constvars]
-  jaxpr = jaxpr.replace(constvars=constvars)
+  jaxpr = jaxpr.replace(constvars=constvars, consts=consts)
   effects = pe.make_jaxpr_effects(jaxpr.constvars, jaxpr.invars,
                                   jaxpr.outvars, jaxpr.eqns)
   jaxpr = jaxpr.replace(effects=effects)

jax/interpreters/partial_eval.py

@@ -48,7 +48,6 @@
   config as config,
   const_fold_rules as const_fold_rules,
   convert_constvars_jaxpr as convert_constvars_jaxpr,
-  convert_envvars_to_constvars as convert_envvars_to_constvars,
   convert_invars_to_constvars as convert_invars_to_constvars,
   custom_partial_eval_rules as custom_partial_eval_rules,
   custom_staging_rules as custom_staging_rules,

tests/core_test.py

@@ -402,13 +402,13 @@ def setUp(self):
     super().setUp()
     lax_control_flow._initial_style_open_jaxpr.cache_clear()
     lax_control_flow._initial_style_jaxpr.cache_clear()
-    lax_control_flow.common._pad_jaxpr_constvars.cache_clear()
+    # lax_control_flow.common._pad_jaxpr_constvars.cache_clear()  # DO_NOT_SUBMIT
 
   def tearDown(self):
     super().tearDown()
     lax_control_flow._initial_style_open_jaxpr.cache_clear()
     lax_control_flow._initial_style_jaxpr.cache_clear()
-    lax_control_flow.common._pad_jaxpr_constvars.cache_clear()
+    # lax_control_flow.common._pad_jaxpr_constvars.cache_clear()  # DO_NOT_SUBMIT
 
   def test_check_jaxpr_correct(self):
     jaxpr = make_jaxpr(lambda x: jnp.sin(x) + jnp.cos(x))(1.).jaxpr

tests/lax_control_flow_test.py

@@ -177,7 +177,7 @@ def setUp(self):
     super().setUp()
     lax_control_flow._initial_style_open_jaxpr.cache_clear()
     lax_control_flow._initial_style_jaxpr.cache_clear()
-    lax_control_flow.common._pad_jaxpr_constvars.cache_clear()
+    # lax_control_flow.common._pad_jaxpr_constvars.cache_clear()  # DO_NOT_SUBMIT
 
   def testCallableErrors(self):
     not_callable = 42