# Copyright 2018 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations from collections import namedtuple from collections.abc import Callable, Sequence import contextlib from dataclasses import dataclass from functools import partial import itertools as it import logging import operator as op from typing import Any, NamedTuple from weakref import ReferenceType, WeakValueDictionary, finalize, ref import numpy as np from jax._src import ad_util from jax._src import api_util from jax._src import config from jax._src import core from jax._src import dtypes from jax._src import effects from jax._src import linear_util as lu from jax._src import profiler from jax._src import source_info_util from jax._src import tree_util from jax._src.core import ( Trace, Tracer, TraceTag, Jaxpr, Literal, AbstractValue, new_jaxpr_eqn, Var, DropVar, Atom, JaxprEqn, Primitive, get_referent, JaxprEqnContext, typeof) from jax._src.lib import _jax from jax._src.source_info_util import SourceInfo from jax._src.state.types import AbstractRef, ReadEffect, RefEffect from jax._src import flattree as ft from jax._src.tree_util import PyTreeDef from jax._src.util import (unzip2, safe_zip, safe_map, toposort, split_list, merge_lists, partition_list, OrderedSet, weakref_lru_cache, multi_weakref_lru_cache, foreach, test_event) map, unsafe_map = safe_map, map zip, unsafe_zip = safe_zip, zip def identity(x): return x TracerId = int AvalId = int ConstId = int AttrKind = Any PyTree = Any HTLV = Any logger = logging.getLogger(__name__) TracebackScope = _jax.TracebackScope class PartialVal(tuple): """Partial value: either a known value or an unknown (abstract) value. Represented as a pair `(aval_opt, const)` of one of two kinds: * `(None, )` indicates a known value, where the constant satisfies `core.valid_jaxtype(const)`; * `(, None)` indicates an unknown value characterized by an abstract value. """ def __new__(cls, xs: tuple[AbstractValue | None, core.Value]): pv, const = xs if config.enable_checks.value: # type checks assert isinstance(pv, (AbstractValue, type(None))), xs assert (const is None or core.valid_jaxtype(const)), const # invariant checks assert (pv is None) ^ (const is None) return tuple.__new__(cls, xs) @classmethod def known(cls, const: core.Value) -> PartialVal: return PartialVal((None, const)) @classmethod def unknown(cls, aval: AbstractValue) -> PartialVal: return PartialVal((aval, None)) def is_known(self) -> bool: return self[0] is None def get_known(self) -> core.Value | None: """Get the known value, if known, else None.""" return self[1] if self[0] is None else None def get_aval(self) -> AbstractValue: """Get AbstractValue directly (if unknown) or from the constant (known).""" known = self.get_known() if known is not None: return typeof(known) else: return self[0] @dataclass(frozen=True, slots=True) class EffectHandle: parents : list[Tracer] recipe : JaxprEqnRecipe class JaxprTrace(Trace): def __init__(self, parent_trace:Trace, name_stack: source_info_util.NameStack, tag:TraceTag): super().__init__() self.name_stack = name_stack self.tag = tag self.parent_trace = parent_trace self.requires_low = False self.effect_handles : list[EffectHandle] = [] self.counter = it.count() def to_jaxpr_tracer(self, x): if isinstance(x, JaxprTracer) and x._trace.tag is self.tag: if x._trace is self: return x else: return JaxprTracer(self, x.pval, FreeVar(x)) else: return self.new_const(x) def stage_value(self, val): return self.to_jaxpr_tracer(val) def new_const(self, val) -> JaxprTracer: return JaxprTracer(self, PartialVal.known(val), None) def new_instantiated_literal(self, val) -> JaxprTracer: aval = typeof(val) return JaxprTracer(self, PartialVal.unknown(aval), Literal(val, aval)) def new_instantiated_const(self, val) -> JaxprTracer: aval = typeof(val) return JaxprTracer(self, PartialVal.unknown(aval), ConstVar(val)) def new_arg(self, pval: PartialVal) -> JaxprTracer: const = pval.get_known() # XXX: Think twice before changing this constant argument pruning! # This has really important consequences for partial_eval_jaxpr. # Most importantly, this guarantees that the unknown jaxpr never uses # known inputs (if it needs them, then they get passed through residuals). if const is None: aval = pval.get_aval() return JaxprTracer(self, PartialVal.unknown(aval), LambdaBinding()) else: return self.new_const(const) def instantiate_const(self, tracer: JaxprTracer) -> JaxprTracer: const = tracer.pval.get_known() if const is None: return tracer else: if core.is_literalable(const, True): return self.new_instantiated_literal(const) else: return self.new_instantiated_const(const) def process_primitive(self, primitive, tracers, params, /): with core.set_current_trace(self.parent_trace): if primitive in custom_partial_eval_rules: tracers = map(self.to_jaxpr_tracer, tracers) return custom_partial_eval_rules[primitive](self, *tracers, **params) else: return self.default_process_primitive(primitive, tracers, params) def default_process_primitive(self, primitive, tracers, params): # By default, if all the input tracers are known, then bind the primitive # and consider all outputs known. Otherwise, stage the application into the # jaxpr and consider all outputs unknown. tracers = map(self.to_jaxpr_tracer, tracers) consts = [t.pval.get_known() for t in tracers] if all(c is not None for c in consts): return primitive.bind_with_trace(self.parent_trace, consts, tuple(t.aval for t in tracers), params) tracers = map(self.instantiate_const, tracers) avals = [t.aval for t in tracers] out_aval, effs = primitive.abstract_eval(*avals, **params) name_stack = self._current_truncated_name_stack() source = source_info_util.current().replace(name_stack=name_stack) if primitive.multiple_results: out_tracers = [JaxprTracer(self, PartialVal.unknown(aval), None) for aval in out_aval] eqn = new_eqn_recipe(self, tracers, out_tracers, primitive, params, effs, source) if effects.partial_eval_kept_effects.filter_in(effs): self.effect_handles.append(EffectHandle(tracers, eqn)) # pyrefly: ignore[bad-argument-type] # pyrefly#2385 for t in out_tracers: t.recipe = eqn return out_tracers else: out_tracer = JaxprTracer(self, PartialVal.unknown(out_aval), None) eqn = new_eqn_recipe(self, tracers, [out_tracer], primitive, params, effs, source) if effects.partial_eval_kept_effects.filter_in(effs): self.effect_handles.append(EffectHandle(tracers, eqn)) # pyrefly: ignore[bad-argument-type] # pyrefly#2385 out_tracer.recipe = eqn return out_tracer def _current_truncated_name_stack(self): return source_info_util.current_name_stack()[len(self.name_stack):] def process_custom_jvp_call(self, prim, fun, jvp, tracers, /, *, symbolic_zeros): tracers = map(self.to_jaxpr_tracer, tracers) if all(t.is_known() for t in tracers): with core.set_current_trace(self.parent_trace): vals = [t.pval[1] for t in tracers] return prim.bind(*vals, subfuns=(fun, jvp), symbolic_zeros=symbolic_zeros) tracers = map(self.instantiate_const, tracers) in_knowns = (False,) * len(tracers) in_avals = tuple(t.aval for t in tracers) in_tangent_avals = [a.to_tangent_aval() for a in in_avals] fun_ = trace_to_subjaxpr_nounits2(fun, self.tag, fun.debug_info, True) fun_, aux = partial_eval_wrapper_nounits(fun_, in_knowns, in_avals) with core.set_current_trace(self.parent_trace): res = fun_.call_wrapped() out_knowns, out_avals, jaxpr, env = aux() assert not any(out_knowns) res_tracers = map(self.instantiate_const, map(self.new_const, res)) env_tracers = map(self.to_jaxpr_tracer, env) out_tracers = [JaxprTracer(self, PartialVal.unknown(a), None) for a in out_avals] closed_jaxpr = convert_constvars_jaxpr(jaxpr) @partial(lu.wrap_init, debug_info=jvp.debug_info) @_memoize def jvp_jaxpr_thunk(*in_zeros): for store in jvp.stores: store and store.reset() nz_tangent_avals, zero_avals = partition_list(in_zeros, in_tangent_avals) jvp_, out_zeros = _jvp_jaxpr_zeros(jvp, in_zeros, tuple(zero_avals)) jvp_jaxpr, _, consts = trace_to_jaxpr_dynamic( jvp_.with_unknown_names(), (*in_avals, *nz_tangent_avals)) return jvp_jaxpr, consts, out_zeros() name_stack = self._current_truncated_name_stack() source = source_info_util.current().replace(name_stack=name_stack) params = dict(call_jaxpr=closed_jaxpr, jvp_jaxpr_fun=jvp_jaxpr_thunk, num_consts=len(res) + len(env), symbolic_zeros=symbolic_zeros) in_tracers = [*res_tracers, *env_tracers, *tracers] effs = core.positional_effects(closed_jaxpr) eqn = new_eqn_recipe(self, in_tracers, out_tracers, prim, params, effs, source) if effects.partial_eval_kept_effects.filter_in(effs): self.effect_handles.append(EffectHandle(in_tracers, eqn)) # pyrefly: ignore[bad-argument-type] # pyrefly#2385 for t in out_tracers: t.recipe = eqn return out_tracers def process_custom_vjp_call(self, prim, f, fwd, bwd, tracers, /, *, out_trees, symbolic_zeros): tracers = map(self.to_jaxpr_tracer, tracers) if all(t.is_known() for t in tracers): vals = [t.pval[1] for t in tracers] with core.set_current_trace(self.parent_trace): return prim.bind(*vals, subfuns=(f, fwd, bwd), out_trees=out_trees, symbolic_zeros=symbolic_zeros) tracers = map(self.instantiate_const, tracers) in_knowns = (False,) * len(tracers) in_avals = tuple(t.aval for t in tracers) f_ = trace_to_subjaxpr_nounits2(f, self.tag, f.debug_info, True) f_, aux = partial_eval_wrapper_nounits(f_, in_knowns, in_avals) params = dict(subfuns=(f_, fwd, bwd), out_trees=out_trees, symbolic_zeros=symbolic_zeros) res = prim.bind_with_trace(self.parent_trace, (), (), params) out_knowns, out_avals, jaxpr, env = aux() assert not any(out_knowns) res_tracers = map(self.instantiate_const, map(self.new_const, res)) env_tracers = map(self.to_jaxpr_tracer, env) out_tracers = [JaxprTracer(self, PartialVal.unknown(a), None) for a in out_avals] closed_jaxpr = convert_constvars_jaxpr(jaxpr) @partial(lu.wrap_init, debug_info=fwd.debug_info) @_memoize def fwd_jaxpr_thunk(*zeros): fwd_ = _interleave_fun(fwd.with_unknown_names(), zeros) fwd_jaxpr, _, consts = trace_to_jaxpr_dynamic(fwd_, in_avals) return fwd_jaxpr, consts name_stack = self._current_truncated_name_stack() source = source_info_util.current().replace(name_stack=name_stack) params = dict( call_jaxpr=closed_jaxpr, fwd_jaxpr_thunk=fwd_jaxpr_thunk, num_consts=len(res) + len(env), bwd=bwd, out_trees=out_trees, symbolic_zeros=symbolic_zeros ) eqn = new_eqn_recipe(self, (*res_tracers, *env_tracers, *tracers), out_tracers, prim, params, core.positional_effects(closed_jaxpr), source) for t in out_tracers: t.recipe = eqn return out_tracers def partition_pvals( pvals: list[PartialVal] ) -> tuple[list[bool], list[AbstractValue], list[Any]]: knowns = [pval.is_known() for pval in pvals ] avals = [pval.get_aval() for pval in pvals if not pval.is_known()] consts = [pval.get_known() for pval in pvals if pval.is_known()] return knowns, avals, consts @lu.transformation_with_aux2 def partial_eval_wrapper_nounits( f: Callable, store: lu.Store, in_knowns: Sequence[bool], in_avals: Sequence[AbstractValue], *in_consts: Any): in_avals_, in_consts_ = iter(in_avals), iter(in_consts) in_pvals = [PartialVal.known(next(in_consts_)) if known else PartialVal.unknown(next(in_avals_)) for known in in_knowns] sentinel = object() assert next(in_avals_, sentinel) is next(in_consts_, sentinel) is sentinel jaxpr, (*maybe_fwds, out_pvals, res, env) = f(in_pvals) out_knowns, out_avals, out_consts = partition_pvals(out_pvals) store.store((*maybe_fwds, out_knowns, out_avals, jaxpr, env)) return (*out_consts, *res) custom_partial_eval_rules: dict[Primitive, Callable] = {} def abstract_eval_fun(fun: Callable, *avals, debug_info: core.DebugInfo, **params): _, avals_out = trace_to_jaxpr( partial(fun, **params), ft.flatten_args(*avals), debug_info) assert all(isinstance(aval, AbstractValue) for aval in avals_out) return list(avals_out) class JaxprTracer(Tracer[JaxprTrace]): __slots__ = ['pval', 'recipe'] _trace: JaxprTrace def __init__(self, trace: JaxprTrace, pval: PartialVal, recipe: JaxprTracerRecipe | None): assert isinstance(pval, PartialVal) super().__init__(trace, pval.get_aval()) self.pval = pval self.recipe = recipe def __repr__(self): return f'Traced<{self.aval}:{self._trace}>' @property def parents(self) -> Sequence[JaxprTracer]: if isinstance(self.recipe, JaxprEqnRecipe): # TODO broadcast_in_dim can create a new tracer... return self.recipe.in_tracers else: return [] def full_lower(self): known = self.pval.get_known() if known is not None: return core.full_lower(known) else: return self def is_known(self): return self.pval.is_known() def get_referent(self): if self.pval.is_known(): return get_referent(self.pval.get_known()) elif isinstance(self.recipe, (FreeVar, ConstVar, Literal)): return get_referent(self.recipe.val) else: return self @profiler.annotate_function def trace_to_jaxpr_nounits( fun: lu.WrappedFun, pvals: Sequence[PartialVal], instantiate: bool | Sequence[bool] = False, ) -> tuple[Jaxpr, list[PartialVal], list[core.Value]]: current_name_stack = source_info_util.current_name_stack() with core.take_current_trace() as parent_trace: trace = JaxprTrace(parent_trace, current_name_stack, TraceTag()) with core.ensure_no_leaks(trace): fun = trace_to_subjaxpr_nounits(fun, trace, instantiate, fun.debug_info) with core.set_current_trace(trace): jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals) assert not env del trace, fun return jaxpr, out_pvals, consts # TODO(mattjj): superfluous wrapper...? @lu.transformation2 def trace_to_subjaxpr_nounits( f: Callable, trace: JaxprTrace, instantiate: Sequence[bool] | bool, debug_info: core.DebugInfo, in_pvals: Sequence[PartialVal]): assert all(isinstance(pv, PartialVal) for pv in in_pvals), in_pvals out_tracers, jaxpr, out_consts, env = _trace_to_subjaxpr_nounits( f, trace, instantiate, in_pvals, debug_info) out_pvals = [t.pval for t in out_tracers] del out_tracers return jaxpr, (out_pvals, out_consts, env) @lu.transformation2 def trace_to_subjaxpr_nounits2( f: Callable, tag: TraceTag, debug_info: core.DebugInfo, instantiate: bool | Sequence[bool], in_pvals: Sequence[PartialVal]): assert isinstance(tag, TraceTag) assert all(isinstance(pv, PartialVal) for pv in in_pvals), in_pvals current_name_stack = source_info_util.current_name_stack() with core.take_current_trace() as parent_trace: trace = JaxprTrace(parent_trace, current_name_stack, tag) out_tracers, jaxpr, out_consts, env = _trace_to_subjaxpr_nounits( f, trace, instantiate, in_pvals, debug_info) out_pvals = [t.pval for t in out_tracers] del out_tracers return jaxpr, (out_pvals, out_consts, env) def _trace_to_subjaxpr_nounits_no_lu_2(f: Callable, trace: JaxprTrace, instantiate: Sequence[bool] | bool, in_pvals: Sequence[PartialVal], debug_info: core.DebugInfo): in_knowns = [pval.is_known() for pval in in_pvals] in_consts = [pval.get_known() for pval in in_pvals if pval.is_known()] in_tracers = [trace.new_arg(pval) for pval in in_pvals if not pval.is_known()] in_args = merge_lists(in_knowns, in_tracers, in_consts) with core.set_current_trace(trace): ans = f(*in_args) assert isinstance(ans, ft.FlatTree), ( f"Got unexpected return type when tracing function to jaxpr: {ans}") assert all(isinstance(x, Tracer) or core.valid_jaxtype(x) for x in ans), ( f"Got unexpected return type when tracing function to jaxpr: {ans}") if isinstance(instantiate, bool): instantiate = [instantiate] * len(ans) out_tracers = ans.map(trace.to_jaxpr_tracer) out_tracers = out_tracers.map2( instantiate, lambda t, inst: trace.instantiate_const(t) if inst else t) out_tracers_ = [t for t in out_tracers if not t.is_known()] jaxpr, out_consts, env = tracers_to_jaxpr( in_tracers, out_tracers_, trace.effect_handles, debug_info.with_unknown_names()) return out_tracers, jaxpr, out_consts, env def _trace_to_subjaxpr_nounits(f: Callable, trace: JaxprTrace, instantiate: Sequence[bool] | bool, in_pvals: Sequence[PartialVal], debug_info: core.DebugInfo): in_knowns = [pval.is_known() for pval in in_pvals] in_consts = [pval.get_known() for pval in in_pvals if pval.is_known()] in_tracers = [trace.new_arg(pval) for pval in in_pvals if not pval.is_known()] in_args = merge_lists(in_knowns, in_tracers, in_consts) with core.set_current_trace(trace): ans = f(*in_args) assert isinstance(ans, (list, tuple)), ( f"Got unexpected return type when tracing function to jaxpr: {ans}") assert all(core.valid_jaxtype(x) for x in ans), ( f"Got unexpected return type when tracing function to jaxpr: {ans}") if isinstance(instantiate, bool): instantiate = [instantiate] * len(ans) out_tracers = map(trace.to_jaxpr_tracer, ans) out_tracers = [trace.instantiate_const(t) if inst else t for inst, t in zip(instantiate, out_tracers)] out_tracers_ = [t for t in out_tracers if not t.is_known()] jaxpr, out_consts, env = tracers_to_jaxpr( in_tracers, out_tracers_, trace.effect_handles, debug_info.with_unknown_names()) return out_tracers, jaxpr, out_consts, env # The below variant implements an optimization where residuals which are also # inputs are indicated in auxiliary data rather than passed as outputs. # TODO(mattjj): update all callers to use this version, delete other version. @lu.transformation2 def trace_to_subjaxpr_nounits_fwd( f: Callable, tag: TraceTag, debug_info: core.DebugInfo, instantiate: bool | Sequence[bool], in_pvals: Sequence[PartialVal]): assert all(isinstance(pv, PartialVal) for pv in in_pvals), in_pvals current_name_stack = source_info_util.current_name_stack() with core.take_current_trace() as parent_trace: trace = JaxprTrace(parent_trace, current_name_stack, tag) with core.set_current_trace(trace): out_tracers, jaxpr, out_consts, env = _trace_to_subjaxpr_nounits( f, trace, instantiate, in_pvals, debug_info) out_pvals = [t.pval for t in out_tracers] # Which out_consts (aka residuals) are just forwarded inputs? Check obj id. in_consts = [pval.get_known() for pval in in_pvals if pval.is_known()] id_map = {id(c): i for i, c in enumerate(in_consts)} fwds: list[int | None] = [id_map.get(id(c)) for c in out_consts] pruned_consts = [c for c, fwd in zip(out_consts, fwds) if fwd is None] del out_tracers return jaxpr, (fwds, out_pvals, pruned_consts, env) # The below variant implements two optimizations: # 1. residuals that are also primal inputs are indicated in aux data rather # than passed as outputs; # 2. residuals that are also primal outputs are indicated in aux data rather # than passed as redundant outputs. def trace_to_subjaxpr_nounits_fwd2( f: Callable, tag: TraceTag, debug_info: core.DebugInfo, instantiate: bool | Sequence[bool], in_pvals: Sequence[PartialVal]): assert all(isinstance(pv, PartialVal) for pv in in_pvals), in_pvals current_name_stack = source_info_util.current_name_stack() with core.take_current_trace() as parent_trace: trace = JaxprTrace(parent_trace, current_name_stack, tag) out_tracers, jaxpr, consts, env = _trace_to_subjaxpr_nounits_no_lu_2( f, trace, instantiate, in_pvals, debug_info) out_pvals = out_tracers.map(lambda t: t.pval) # Which consts (aka residuals) are just forwarded inputs? Check obj id. in_consts = [pval.get_known() for pval in in_pvals if pval.is_known()] id_map = {id(c): i for i, c in enumerate(in_consts)} input_fwds: list[int | None] = [id_map.get(id(c)) for c in consts] # Which consts (aka residuals) are already primal outputs? Check obj id. out_consts = [pval.get_known() for pval in out_pvals if pval.is_known()] id_map = {id(c): i for i, c in enumerate(out_consts)} output_fwds: list[int | None] = [id_map.get(id(c)) for c in consts] pruned_consts = [c for c, f1, f2 in zip(consts, input_fwds, output_fwds) if f1 is None and f2 is None] del out_tracers return jaxpr, (input_fwds, output_fwds, out_pvals, pruned_consts, env) FreeVar = namedtuple('FreeVar', ['val']) ConstVar = namedtuple('ConstVar', ['val']) LambdaBinding = namedtuple('LambdaBinding', []) class JaxprEqnRecipe(NamedTuple): eqn_id: Any in_tracers: Sequence[JaxprTracer] out_tracer_refs: Sequence[ref[JaxprTracer]] out_avals: Sequence[core.AbstractValue] primitive: Primitive params: dict[str, Any] effects: core.Effects source_info: source_info_util.SourceInfo ctx: JaxprEqnContext JaxprTracerRecipe = ( JaxprEqnRecipe | LambdaBinding | FreeVar | ConstVar | Literal ) def new_eqn_recipe(trace: JaxprTrace, in_tracers: Sequence[JaxprTracer], out_tracers: Sequence[JaxprTracer], primitive: Primitive, params: dict[str, Any], effects: core.Effects, source_info: source_info_util.SourceInfo, ctx: JaxprEqnContext | None = None) -> JaxprEqnRecipe: out_avals = [t.aval for t in out_tracers] ctx = ctx or core.current_jaxpr_eqn_context() return JaxprEqnRecipe(next(trace.counter), tuple(in_tracers), map(ref, out_tracers), out_avals, primitive, params, effects, source_info, ctx) def tracers_to_jaxpr( in_tracers: Sequence[JaxprTracer], out_tracers: Sequence[JaxprTracer], effect_handles: Sequence[Any], debug_info: core.DebugInfo, ) -> tuple[Jaxpr, tuple[Any, ...], tuple[Any, ...]]: """Constructs Jaxpr given tracers for inputs and outputs. Params: in_tracers: the tracers that were created for the function inputs out_tracers: the tracers that were output by the function. debug_info: the debug info for the function. Returns: a triple of a `Jaxpr`, a list of constant values corresponding to the `constvars` in the returned Jaxps, and a list of environment values. The vars for the environment values have been prepended to the Jaxpr's `invars`. """ gensym = core.gensym() t_to_var: dict[TracerId, Var] = {} consts: dict[Var, Any] = {} env: dict[Var, JaxprTracer] = {} constid_to_var: dict[ConstId, Var] = {} # for deduplication def get_atom(t: JaxprTracer) -> Atom: return t.recipe if type(t.recipe) is Literal else t_to_var[id(t)] def newvar(t: JaxprTracer | None) -> Var: assert t is not None var = gensym(t.aval) var_ = t_to_var.setdefault(id(t), var) assert var is var_ return var processed_eqn_ids = set() eqns: list[core.JaxprEqn] = [] is_high = False reachable = toposort tracers = reachable((*in_tracers, *out_tracers, *effect_handles)) def sort_key(t): r = t.recipe return r.eqn_id if isinstance(r, JaxprEqnRecipe) else -1 tracers = sorted(tracers, key=sort_key) for t in tracers: r = t.recipe if isinstance(r, JaxprEqnRecipe): # TODO broadcast_in_dim can create a new tracer, not present in parents if r.eqn_id not in processed_eqn_ids: in_atoms = map(get_atom, r.in_tracers) outvars = [DropVar(a) if rf() is None else newvar(rf()) for a, rf in zip(r.out_avals, r.out_tracer_refs)] eqns.append(new_jaxpr_eqn(in_atoms, outvars, r.primitive, r.params, r.effects, r.source_info, r.ctx)) in_avals = [x.aval for x in in_atoms] is_high |= r.primitive.is_high(*in_avals, **r.params) processed_eqn_ids.add(r.eqn_id) elif isinstance(r, LambdaBinding): if not any(t is in_tracer for in_tracer in in_tracers): raise core.escaped_tracer_error(t, f"Tracer not in input tracers: {t}") newvar(t) elif isinstance(r, ConstVar): var = constid_to_var.get(id(r.val)) if var is None: var = constid_to_var[id(r.val)] = newvar(t) consts[var] = r.val t_to_var[id(t)] = var elif isinstance(r, FreeVar): env[newvar(t)] = r.val elif isinstance(r, Literal): pass elif r is None: assert False else: raise TypeError(r) env_vars, env_vals = unzip2(env.items()) invars = [*env_vars, *map(get_atom, in_tracers)] const_vars, const_vals = unzip2(consts.items()) outvars = map(get_atom, out_tracers) jaxpr_effects = make_jaxpr_effects(const_vars, invars, outvars, eqns) is_high |= any(x.aval.is_high for x in it.chain(const_vars, invars, outvars)) jaxpr = Jaxpr(const_vars, invars, # pyrefly: ignore[bad-argument-type] outvars, eqns, jaxpr_effects, debug_info, is_high) config.enable_checks.value and core.check_jaxpr(jaxpr) # del getvar # needed to avoid cyclic-reference closure, apparently! return jaxpr, const_vals, env_vals @weakref_lru_cache def move_envvars(jaxpr: Jaxpr, which: tuple[bool, ...]) -> Jaxpr: """Move the leading invars selected by `which` after the unselected ones.""" assert not jaxpr.consts keep, env = partition_list(which, jaxpr.invars[:len(which)]) return jaxpr.replace(invars=[*keep, *env, *jaxpr.invars[len(which):]]) def separate_consts(jaxpr: Jaxpr) -> tuple[Jaxpr, list[Any]]: """Detaches the consts and returns them explicitly.""" return convert_constvars_jaxpr(jaxpr), jaxpr.consts def convert_constvars_jaxpr(jaxpr: Jaxpr) -> Jaxpr: """Detaches the consts, exposing the constant inputs as leading invars.""" return _detach_consts(jaxpr) if jaxpr.consts else jaxpr @weakref_lru_cache def _detach_consts(jaxpr: Jaxpr) -> Jaxpr: return jaxpr.with_consts(()) def partial_eval_jaxpr_nounits( jaxpr: Jaxpr, unknowns: Sequence[bool], instantiate: bool | Sequence[bool], ) -> tuple[Jaxpr, Jaxpr, list[bool], list[AbstractValue]]: """Unzip a jaxpr in two by data dependence into 'known' and 'unknown' parts. That is, given a jaxpr and a sequence of booleans indicating which jaxpr inputs (i.e. invars) are considered unknown, produce two jaxprs, a list of booleans representing which of the original jaxpr's outputs are unknown (i.e. have a data dependence on an unknown input), and a list of abstract values representing residuals (part of the first jaxpr's output and the second jaxpr's input). The two jaxprs result from partitioning the original jaxpr's first-order primitive applications based on whether all the inputs to the application are known (in which case the application is represented in the 'known' jaxpr and its result is considered known) or whether any inputs to the application are unknown (in which case the application is represented in the 'unknown' jaxpr and its result is considered unknown). Higher-order primitives are recursively unzipped in two. The `instantiate` argument can be used to ensure some outputs are lifted into the 'unknown' jaxpr. For example, give an input jaxpr: { lambda ; a:f32[] b:f32[]. let c:f32[] = cos a d:f32[] = sin a e:f32[] = neg d f:f32[] = mul e b in (c, f) } then applying this function with `unknowns=[False, True]` and `instantiate=False` produces as an output triple: # jaxpr_known { lambda ; a:f32[]. let b:f32[] = cos a c:f32[] = sin a d:f32[] = neg c in (b, d) } # jaxpr_unknown { lambda ; a:f32[] b:f32[]. let c:f32[] = mul b a in (c,) } # out_unknowns [False, True] Notice in particular that the first output (jaxpr_known) contains all the primitive applications which do not have a data dependence on an unknown input. Also notice the input and output types: the input type of the first jaxpr produced represents the type of the known inputs of the original jaxpr, and the output type of the second jaxpr produced represents the type of the unknown outputs of the original jaxpr. In the above example, the output of jaxpr_known named `d` is a _residual_ output, and corresponds to the input named `a` in jaxpr_unknown. In general, jaxpr_known will produce extra outputs (at the end of its output list) corresponding to intermediate values of the original jaxpr which must be passed to jaxpr_unknown (as leading inputs). """ instantiate = tuple(instantiate) if isinstance(instantiate, list) else instantiate return _partial_eval_jaxpr_nounits(jaxpr, tuple(unknowns), instantiate, False)[:-1] def partial_eval_jaxpr_nounits_fwd( jaxpr: Jaxpr, unknowns: Sequence[bool], instantiate: bool | Sequence[bool], fwd: bool | Sequence[bool] = True, ) -> tuple[Jaxpr, Jaxpr, list[bool], list[AbstractValue], list[int | None]]: instantiate = tuple(instantiate) if isinstance(instantiate, list) else instantiate fwd = tuple(fwd) if isinstance(fwd, list) else fwd return _partial_eval_jaxpr_nounits(jaxpr, tuple(unknowns), instantiate, fwd) @weakref_lru_cache def _partial_eval_jaxpr_nounits( jaxpr: Jaxpr, in_unknowns: Sequence[bool], instantiate: bool | Sequence[bool], fwd: bool | Sequence[bool]): f = lu.wrap_init(core.jaxpr_as_fun(jaxpr), debug_info=jaxpr.debug_info) cell = [] def fun(*known_vals_in): known_vals_in_ = iter(known_vals_in) unknown_avals = (a for a, uk in zip(jaxpr.in_avals, in_unknowns) if uk) in_pvals = [PartialVal.unknown(next(unknown_avals)) if uk else PartialVal.known(next(known_vals_in_)) for uk in in_unknowns] assert next(known_vals_in_, None) is next(unknown_avals, None) is None jaxpr_unknown_, (fwds, out_pvals, residuals, ()) = trace_to_subjaxpr_nounits_fwd( f, TraceTag(), jaxpr.debug_info, instantiate).call_wrapped(in_pvals) jaxpr_unknown = convert_constvars_jaxpr(jaxpr_unknown_) out_unknowns = [not pval.is_known() for pval in out_pvals] if type(fwd) is bool and not fwd: residuals_ = iter(residuals) residuals = [next(residuals_) if f is None else known_vals_in[f] for f in fwds] assert next(residuals_, None) is None fwds = [None] * len(fwds) else: if type(fwd) is tuple: fwd_ = [f for f, uk in zip(fwd, in_unknowns) if not uk] residuals_, residuals = iter(residuals), [] fwds = [residuals.append(next(residuals_)) if f is None else residuals.append(known_vals_in[f]) if not fwd_[f] else f for f in fwds] fwds, residuals = _include_consts_in_fwds(jaxpr.consts, fwds, residuals) res_avals = [core.typeof(r) for r in residuals] cell.append((out_unknowns, jaxpr_unknown, res_avals, fwds)) known_vals_out = [pval.get_known() for pval in out_pvals if pval.is_known()] return [*known_vals_out, *residuals] known_avals = [a for a, uk in zip(jaxpr.in_avals, in_unknowns) if not uk] closed_jaxpr_known, _ = trace_to_jaxpr( fun, ft.flatten_args(*known_avals), f.debug_info.with_unknown_names()) (out_unknowns, jaxpr_unknown, res_avals, fwds), = cell if config.enable_checks.value: core.check_jaxpr(closed_jaxpr_known) core.check_jaxpr(jaxpr_unknown) closed_jaxpr_unknown = jaxpr_unknown return closed_jaxpr_known, closed_jaxpr_unknown, out_unknowns, res_avals, fwds def _include_consts_in_fwds(consts, fwds, residuals): if all(f is None for f in fwds): return fwds, residuals dummys = [object() for _ in range(max(f for f in fwds if f is not None) + 1)] residuals_ = iter(residuals) residuals = [next(residuals_) if f is None else dummys[f] for f in fwds] assert next(residuals_, None) is None idxs = {id(x): i for i, x in enumerate((*consts, *dummys))} fwds = [idxs.get(id(r)) for r in residuals] residuals = [r for r in residuals if id(r) not in idxs] return fwds, residuals def partial_eval_jaxpr_custom( jaxpr: Jaxpr, in_unknowns: Sequence[bool], in_inst: bool | Sequence[bool], ensure_out_unknowns: bool | Sequence[bool], ensure_out_inst: bool | Sequence[bool], saveable: Callable[..., RematCases_], ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int]: *outs, num_res_ref = partial_eval_jaxpr_stateful( jaxpr, in_unknowns, in_inst, ensure_out_unknowns, ensure_out_inst, saveable) if num_res_ref: raise ValueError("Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.") return *outs, # pyrefly: ignore[bad-return] def partial_eval_jaxpr_stateful( jaxpr: Jaxpr, in_unknowns: Sequence[bool], in_inst: bool | Sequence[bool], ensure_out_unknowns: bool | Sequence[bool], ensure_out_inst: bool | Sequence[bool], saveable: Callable[..., RematCases_] | None, ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int, int]: if type(in_inst) is bool: in_inst = (in_inst,) * len(jaxpr.invars) if type(ensure_out_unknowns) is bool: ensure_out_unknowns = (ensure_out_unknowns,) * len(jaxpr.outvars) if type(ensure_out_inst) is bool: ensure_out_inst = (ensure_out_inst,) * len(jaxpr.outvars) if saveable is None: saveable = everything_saveable jaxpr_known, jaxpr_staged, out_unknowns, out_inst, num_res, num_res_ref = \ _partial_eval_jaxpr_custom_cached( jaxpr, tuple(in_unknowns), tuple(in_inst), tuple(ensure_out_unknowns), tuple(ensure_out_inst), saveable) return jaxpr_known, jaxpr_staged, out_unknowns, out_inst, num_res, num_res_ref everything_saveable = lambda *_, **__: True @weakref_lru_cache def _partial_eval_jaxpr_custom_cached( jaxpr: Jaxpr, in_unknowns: tuple[bool, ...], in_inst: tuple[bool, ...], ensure_out_unknowns: tuple[bool, ...], ensure_out_inst: tuple[bool, ...], saveable: Callable[..., RematCases_], ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int, int]: env: dict[Var, tuple[bool, bool]] = {} residuals: OrderedSet[Var] = OrderedSet() residual_refs: OrderedSet[Var] = OrderedSet() def read(x: Atom) -> tuple[bool, bool]: if type(x) is Var: return env[x] return (False, True) def write(unk: bool, inst: bool, v: Var) -> None: assert (unk, inst) != (True, False) env[v] = (unk, inst) def ensure_instantiated(inst: bool, x: Atom) -> Atom: if type(x) is Var and not inst: residuals.add(x) return x def has_effects(effects) -> bool: not_really_effects = (core.NamedAxisEffect, core.InternalMutableArrayEffect) return any(not isinstance(e, not_really_effects) for e in effects) known_eqns, staged_eqns = [], [] foreach(write, in_unknowns, in_inst, jaxpr.invars) foreach(partial(write, False, True), jaxpr.constvars) for eqn in jaxpr.eqns: unks_in, inst_in = unzip2(map(read, eqn.invars)) rule = partial_eval_jaxpr_custom_rules.get(eqn.primitive) if rule: eqn1, eqn2, unks_out, inst_out, res = rule(saveable, unks_in, inst_in, eqn) eqn1 and known_eqns.append(eqn1); eqn2 and staged_eqns.append(eqn2) for r in res: if isinstance(r.aval, AbstractRef): residual_refs.add(r) else: residuals.add(r) foreach(write, unks_out, inst_out, eqn.outvars) elif any(unks_in): inputs = map(ensure_instantiated, inst_in, eqn.invars) staged_eqns.append(eqn.replace(invars=inputs)) foreach(partial(write, True, True), eqn.outvars) else: known_eqns.append(eqn) # If it's an effectful primitive, we always to run and avoid staging it. policy = ensure_enum(saveable( eqn.primitive, *[x.aval for x in eqn.invars], **eqn.params)) if has_effects(eqn.effects) or isinstance(policy, SaveableType): foreach(partial(write, False, False), eqn.outvars) elif isinstance(policy, Offloadable): # TODO(slebedev): This is a legit error which requires a BUILD fix. from jax._src.dispatch import device_put_p, ArrayCopySemantics # pyrefly: ignore[missing-import] resvars = [Var(v.aval.update(memory_space=core.mem_kind_to_space(policy.dst))) for v in eqn.outvars] offload_eqn = core.JaxprEqn( eqn.outvars, resvars, device_put_p, dict( devices=(core.mem_kind_to_space(policy.dst),) * len(eqn.outvars), srcs=(None,), copy_semantics=(ArrayCopySemantics.ALWAYS_COPY,), ), set(), source_info_util.new_source_info(), core.current_jaxpr_eqn_context()) known_eqns.append(offload_eqn) # resvars are known and available in the backward jaxpr. foreach(partial(write, False, True), resvars) assert all(o.aval.memory_space == core.mem_kind_to_space(policy.src) # pyrefly: ignore[missing-attribute] for o in eqn.outvars) residuals.update(resvars) reload_eqn = core.JaxprEqn( resvars, eqn.outvars, device_put_p, dict( devices=(core.mem_kind_to_space(policy.src),) * len(resvars), srcs=(None,), copy_semantics=(ArrayCopySemantics.ALWAYS_COPY,) ), set(), source_info_util.new_source_info(), core.current_jaxpr_eqn_context()) staged_eqns.append(reload_eqn) # outvars are known and available in the backward jaxpr. foreach(partial(write, False, True), eqn.outvars) else: assert isinstance(policy, RecomputeType) inputs = map(ensure_instantiated, inst_in, eqn.invars) staged_eqns.append(eqn.replace(invars=inputs)) foreach(partial(write, False, True), eqn.outvars) unzipped = unzip2(map(read, jaxpr.outvars)) out_unknowns, out_inst = list(unzipped[0]), list(unzipped[1]) assert all(type(v) is Var for v in residuals), residuals for x, inst, ensure_inst in zip(jaxpr.outvars, out_inst, ensure_out_inst): if ensure_inst: ensure_instantiated(inst, x) out_unknowns = map(op.or_, out_unknowns, ensure_out_unknowns) out_inst = map(op.or_, out_inst, ensure_out_inst) ins_known, _ = partition_list(in_unknowns, jaxpr.invars) outs_known, _ = partition_list(out_unknowns, jaxpr.outvars) ref_res_is_input = [r in ins_known for r in residual_refs] non_input_res_refs, _ = partition_list(ref_res_is_input, list(residual_refs)) ins_known_and_ref_res = [*ins_known, *non_input_res_refs] known_outvars = [*outs_known, *residuals] known_effects = make_jaxpr_effects(jaxpr.constvars, ins_known_and_ref_res, known_outvars, known_eqns) # TODO(mattjj,necula): debug info should be updated here jaxpr_known = jaxpr.replace( invars=ins_known_and_ref_res, outvars=known_outvars, eqns=known_eqns, effects=known_effects, debug_info=jaxpr.debug_info.with_unknown_names()) config.enable_checks.value and core.check_jaxpr(jaxpr_known) _, ins_staged = partition_list(in_inst, jaxpr.invars) _, outs_staged = partition_list(out_inst, jaxpr.outvars) staged_invars = [*residuals, *non_input_res_refs, *ins_staged] staged_effects = make_jaxpr_effects(jaxpr.constvars, staged_invars, outs_staged, staged_eqns) # TODO(mattjj,necula): debug info should be updated here jaxpr_staged = jaxpr.replace( invars=staged_invars, outvars=outs_staged, eqns=staged_eqns, effects=staged_effects, debug_info=jaxpr.debug_info.with_unknown_names()) config.enable_checks.value and core.check_jaxpr(jaxpr_staged) return (jaxpr_known, jaxpr_staged, out_unknowns, out_inst, len(residuals), len(non_input_res_refs)) MemoryKind = str class RecomputeType: pass Recompute = RecomputeType() class SaveableType: pass Saveable = SaveableType() class Offloadable(NamedTuple): src: MemoryKind dst: MemoryKind RematCases = RecomputeType | SaveableType | Offloadable RematCases_ = RematCases | bool def ensure_enum(case: bool | RematCases) -> RematCases: if isinstance(case, bool): return Saveable if case else Recompute if not isinstance(case, (RecomputeType, SaveableType, Offloadable)): msg = ("Value returned by a remat policy should be a bool or" " `ad_checkpoint.Recompute`, `ad_checkpoint.Saveable` or" " `ad_checkpoint.Offloadable(...)`." f" Got {case} of type {type(case)}.") if isinstance(case, Offloadable): msg += ("Did you return `Offloadable` instead of an instantiated" " `Offloadable(...)`?") raise TypeError(msg) return case # A primitive rule for policy-driven partial evaluation returns a 5-tuple # with the components representing, respectively: # * the JaxprEqn for the 'known' side (or None if there is no known component), # * the JaxprEqn for the 'unknown' side (or None), # * a list of booleans indicating which of the original outputs are unknown, # * a list of booleans indicating which of the original outputs are # instantiated (i.e. available) in the 'unknown' side, # * a list of Var instances representing residuals to be added (i.e. to be # plumbed as outputs of the 'known' side jaxpr and added as input binders to # the 'unknown' jaxpr). PartialEvalCustomResult = tuple[JaxprEqn | None, JaxprEqn | None, Sequence[bool], Sequence[bool], list[Var]] PartialEvalCustomRule = Callable[ [Callable[..., RematCases_], Sequence[bool], Sequence[bool], JaxprEqn], PartialEvalCustomResult] partial_eval_jaxpr_custom_rules: dict[Primitive, PartialEvalCustomRule] = {} ParamsUpdater = Callable[[Sequence[bool], Sequence[bool], Sequence[bool], Sequence[bool], int, dict, dict], tuple[dict, dict]] ResAvalUpdater = Callable[[dict[str, Any], AbstractValue], AbstractValue] def _default_res_aval_updater( params: dict[str, Any], aval: AbstractValue) -> AbstractValue: return aval def call_partial_eval_custom_rule( jaxpr_param_name: str, params_updater: ParamsUpdater, saveable: Callable[..., RematCases_], unks_in: list[bool], inst_in: list[bool], eqn: JaxprEqn, *, res_aval: ResAvalUpdater = _default_res_aval_updater, ctx = contextlib.nullcontext, ) -> tuple[JaxprEqn, JaxprEqn, Sequence[bool], Sequence[bool], list[Var]]: jaxpr = eqn.params[jaxpr_param_name] with ctx(eqn.params): jaxpr_known, jaxpr_staged, unks_out, inst_out, num_res = \ partial_eval_jaxpr_custom(jaxpr, unks_in, inst_in, False, False, saveable) ins_known, _ = partition_list(unks_in, eqn.invars) out_binders_known, _ = partition_list(unks_out, eqn.outvars) _, ins_staged = partition_list(inst_in, eqn.invars) _, out_binders_staged = partition_list(inst_out, eqn.outvars) params_known = {**eqn.params, jaxpr_param_name: jaxpr_known} params_staged = {**eqn.params, jaxpr_param_name: jaxpr_staged} params_known, params_staged = params_updater( unks_in, inst_in, map(op.not_, unks_out), inst_out, num_res, params_known, params_staged) residuals = [Var(res_aval(params_known, var.aval)) for var in jaxpr_staged.invars[:num_res]] eqn_known = new_jaxpr_eqn( ins_known, [*out_binders_known, *residuals], eqn.primitive, params_known, core.eqn_effects(jaxpr_known, ins_known), eqn.source_info, eqn.ctx) eqn_staged = new_jaxpr_eqn( [*residuals, *ins_staged], out_binders_staged, eqn.primitive, params_staged, core.eqn_effects(jaxpr_staged, [*residuals, *ins_staged]), eqn.source_info, eqn.ctx) assert len(eqn_staged.invars) == len(jaxpr_staged.invars) new_inst = [x for x, inst in zip(eqn.invars, inst_in) if type(x) is Var and not inst] return eqn_known, eqn_staged, unks_out, inst_out, new_inst + residuals # TODO(mattjj): unify with ParamsUpdater (this one takes an extra int) ParamsUpdater2 = Callable[[Sequence[bool], Sequence[bool], Sequence[bool], Sequence[bool], int, int, dict, dict], tuple[dict, dict]] def closed_call_partial_eval_custom_rule( jaxpr_param_name: str, params_updater: ParamsUpdater2, saveable: Callable[..., RematCases_], unks_in: list[bool], inst_in: list[bool], eqn: JaxprEqn, *, res_aval: ResAvalUpdater = _default_res_aval_updater, ) -> tuple[JaxprEqn, JaxprEqn, Sequence[bool], Sequence[bool], list[Var]]: # TODO(sharadmv,mattjj): dedup this rule with call_partial_eval_custom_rule. disallow_output_fwds = tuple(isinstance(v, DropVar) for v in eqn.outvars) # TODO(mattjj): this is just for pjit... but let's delete all this code from jax._src.sharding_impls import UNSPECIFIED # pyrefly: ignore[missing-import] in_shardings, in_layouts = eqn.params.get('in_shardings'), eqn.params.get('in_layouts') if in_shardings is not None: assert in_layouts is not None disallow_input_fwds = tuple(s is not UNSPECIFIED or l is not None for s, l in zip(in_shardings, in_layouts)) else: disallow_input_fwds = (False,) * len(unks_in) jaxpr_known, jaxpr_staged, unks_out, inst_out, num_res_ref, num_res_val, in_fwd, out_fwd = \ _closed_jaxpr_partial_eval_custom_cached( eqn.params[jaxpr_param_name], (*unks_in,), (*inst_in,), disallow_input_fwds, disallow_output_fwds, saveable) num_res = num_res_ref + num_res_val out_binders_known, _ = partition_list(unks_out, eqn.outvars) ins_known, _ = partition_list(unks_in, eqn.invars) _, ins_staged = partition_list(inst_in, eqn.invars) _, out_binders_staged = partition_list(inst_out, eqn.outvars) params_known = {**eqn.params, jaxpr_param_name: jaxpr_known} params_staged = {**eqn.params, jaxpr_param_name: jaxpr_staged} params_known, params_staged = params_updater( unks_in, inst_in, map(op.not_, unks_out), inst_out, sum(fin is fout is None for fin, fout in zip(in_fwd, out_fwd)), num_res, params_known, params_staged) res_val_binders, res_ref_binders = split_list( [Var(res_aval(params_known, v)) for v in jaxpr_staged.in_avals[:num_res]], [num_res_val]) res_val_binders = [v for v, fin, fout in zip(res_val_binders, in_fwd, out_fwd) if fin is fout is None] res_val_binders_ = iter(res_val_binders) res_val_vars = [out_binders_known[fout] if fout is not None else ins_known[fin] if fin is not None else next(res_val_binders_) for fin, fout in zip(in_fwd, out_fwd)] assert next(res_val_binders_, None) is None eqn_known = new_jaxpr_eqn( [*ins_known, *res_ref_binders], [*out_binders_known, *res_val_binders], eqn.primitive, params_known, core.eqn_effects(jaxpr_known, [*ins_known, *res_ref_binders]), eqn.source_info, eqn.ctx) eqn_staged = new_jaxpr_eqn( [*res_val_vars, *res_ref_binders, *ins_staged], out_binders_staged, eqn.primitive, params_staged, core.eqn_effects(jaxpr_staged, [*res_val_vars, *res_ref_binders, *ins_staged]), eqn.source_info, eqn.ctx) assert len(eqn_staged.invars) == len(jaxpr_staged.in_avals) assert len(ins_known) + len(res_ref_binders) == len(jaxpr_known.invars) assert len(ins_staged) + len(res_ref_binders) + len(res_val_vars) == len(jaxpr_staged.invars) assert len(out_binders_known) + len(res_val_binders) == len(jaxpr_known.outvars) new_inst = [x for x, inst in zip(eqn.invars, inst_in) if type(x) is Var and not inst] new_vars = [*new_inst, *res_val_vars, *res_ref_binders] return eqn_known, eqn_staged, unks_out, inst_out, new_vars # pyrefly: ignore[bad-return] @weakref_lru_cache def _closed_jaxpr_partial_eval_custom_cached( jaxpr: Jaxpr, unks_in: tuple[bool, ...], inst_in: tuple[bool, ...], disallowed_input_forwards: tuple[bool, ...], disallowed_output_forwards: tuple[bool, ...], saveable: Callable ) -> tuple[Jaxpr, Jaxpr, Sequence[bool], Sequence[bool], int, int, Sequence[int | None], Sequence[int | None]]: jaxpr_known_, jaxpr_staged_, unks_out, inst_out, num_res_val, num_res_ref = \ partial_eval_jaxpr_stateful(jaxpr, unks_in, inst_in, False, False, saveable) num_out_primals = len(jaxpr_known_.outvars) - num_res_val out_vars, res_vars = split_list(jaxpr_known_.outvars, [num_out_primals]) # Compute which residual value outputs are also primal inputs. disallowed, _ = partition_list(unks_in, disallowed_input_forwards) idx_map = {id(v): i for i, (v, b) in enumerate(zip(jaxpr_known_.invars, disallowed)) if not b} in_fwd = [idx_map.get(id(v)) for v in res_vars] # Compute which residual value outputs are also *undropped* primal outputs. disallowed, _ = partition_list(unks_out, disallowed_output_forwards) idx_map = {id(v): i for i, (v, b) in enumerate(zip(out_vars, disallowed)) if not b} out_fwd = [idx_map.get(id(v)) for v in res_vars] # Prune jaxpr_known_ outputs by removing forwards. keep = [f1 is f2 is None for f1, f2 in zip(in_fwd, out_fwd)] jaxpr_known_ = prune_jaxpr_outputs(jaxpr_known_, [True] * num_out_primals + keep) return (jaxpr_known_, jaxpr_staged_, unks_out, inst_out, num_res_ref, num_res_val, in_fwd, out_fwd) def _jaxpr_forwarding(jaxpr: Jaxpr) -> list[int | None]: # Compute which inputs are just forwarded to outputs. fwds: dict[Var, Atom] = dict(zip(jaxpr.invars, jaxpr.invars)) for eqn in jaxpr.eqns: if eqn.primitive in forwarding_rules: eqn = eqn.replace(invars=[a if type(a) is Literal else fwds.get(a, a) for a in eqn.invars]) fwd_idx, _ = forwarding_rules[eqn.primitive](eqn) for v_orig, idx in zip(eqn.outvars, fwd_idx): if idx is not None: fwds[v_orig] = eqn.invars[idx] idxs: dict[Var, int] = {v: i for i, v in enumerate(jaxpr.invars)} return [None if type(v) is Literal else idxs.get(fwds.get(v)) # pyrefly: ignore[bad-argument-type] for v in jaxpr.outvars] def prune_jaxpr_outputs(jaxpr: Jaxpr, used_outputs: Sequence[bool]) -> Jaxpr: return _prune_jaxpr_outputs_cached(jaxpr, tuple(used_outputs)) def _prune_jaxpr_outputs(jaxpr: Jaxpr, used_outputs: tuple[bool, ...]) -> Jaxpr: outvars = [v for v, b in zip(jaxpr.outvars, used_outputs) if b] dbg = core.DebugInfo( jaxpr.debug_info.traced_for, jaxpr.debug_info.func_src_info, jaxpr.debug_info.arg_names, jaxpr.debug_info.filter_result_paths(used_outputs)) new_jaxpr = jaxpr.replace(outvars=outvars, debug_info=dbg) config.enable_checks.value and core.check_jaxpr(new_jaxpr) return new_jaxpr _prune_jaxpr_outputs_cached = weakref_lru_cache(_prune_jaxpr_outputs) # Since the Jaxpr/Jaxpr merge this is the same as prune_jaxpr_outputs # (attached consts are preserved by Jaxpr.replace). prune_closed_jaxpr_outputs = prune_jaxpr_outputs def dedup_jaxpr_outputs(jaxpr: Jaxpr, num_kept_prefix: int, fwdable_prefix: Sequence[bool] | None = None, ) -> tuple[Jaxpr, list[int | None]]: """Prune duplicated outputs, for callers to restore after evaluation.""" prefix_vars, rest_vars = split_list(jaxpr.outvars, [num_kept_prefix]) fwdable = ([True] * num_kept_prefix if fwdable_prefix is None else fwdable_prefix) idx_map = {id(v): i for i, (v, f) in enumerate(zip(prefix_vars, fwdable)) if f} num_kept = num_kept_prefix out_fwd: list[int | None] = [None] * num_kept_prefix for v in rest_vars: if (fwd := idx_map.get(id(v))) is None: idx_map[id(v)] = num_kept num_kept += 1 out_fwd.append(fwd) jaxpr = prune_jaxpr_outputs(jaxpr, [f is None for f in out_fwd]) return jaxpr, out_fwd def dce_jaxpr(jaxpr: Jaxpr, used_outputs: bool | Sequence[bool], instantiate: bool | Sequence[bool] = False, live_inputs: bool | Sequence[bool] = True, ) -> tuple[Jaxpr, list[bool]]: """Runs dead-code elementation on a given jaxpr. Args: jaxpr: The jaxpr to DCE. used_outputs: A list of bools indicating which outputs are used. instantiate: A bool or a list of bools indicating which inputs should be considered used, regardless of whether they are actually used in a jaxpr. If a bool, the same value is used for all inputs. live_inputs: A bool or a list of bools indicating, for each Ref-typed input, whether its contents may be read after the jaxpr returns (e.g. by the caller). Writes to a Ref input marked False may be removed if nothing in the jaxpr reads them. Entries for non-Ref inputs are ignored. If a bool, the same value is used for all inputs. Returns: A tuple of ``(new_jaxpr, used_inputs)``. """ if type(used_outputs) is bool: used_outputs = (used_outputs,) * len(jaxpr.outvars) if type(instantiate) is bool: instantiate = (instantiate,) * len(jaxpr.invars) if type(live_inputs) is bool: live_inputs = (live_inputs,) * len(jaxpr.invars) new_jaxpr, used_inputs = _dce_jaxpr(jaxpr, tuple(used_outputs), tuple(instantiate), tuple(live_inputs)) return (jaxpr if new_jaxpr is None else new_jaxpr), used_inputs def dce_jaxpr_consts(jaxpr: Jaxpr, used_outputs: Sequence[bool], instantiate: bool | Sequence[bool] = False, ) -> tuple[Jaxpr, list[bool], list[bool]]: new_jaxpr, used_inputs_ = dce_jaxpr(convert_constvars_jaxpr(jaxpr), used_outputs, instantiate) used_consts, used_inputs = split_list(used_inputs_, [len(jaxpr.constvars)]) new_jaxpr = new_jaxpr.with_consts( [c for c, used in zip(jaxpr.consts, used_consts) if used]) return new_jaxpr, used_consts, used_inputs def _default_dce_rule( used_outs: list[bool], live_ins: list[bool], eqn: JaxprEqn ) -> tuple[list[bool], JaxprEqn | None]: if not any(used_outs) and not has_effects(eqn, live_ins): return [False] * len(eqn.invars), None return [True] * len(eqn.invars), eqn dce_rules: dict[Primitive, DCERule] = {} dceable_effects = effects.EffectTypeSet() dceable_effects.add_type(ReadEffect) dceable_effects.add_type(core.NamedAxisEffect) dceable_effects.add_type(core.InternalMutableArrayEffect) def _free_ref_dce_rule( used_outs: list[bool], live_ins: list[bool], eqn: JaxprEqn ) -> tuple[list[bool], JaxprEqn | None]: # Never gonna DCE free_ref. # TODO(mattjj): free_ref doesn't read its Ref, but returning True here marks # the Ref live, so we keep any writes to it that precede the free_ref. del used_outs, live_ins return [True] * len(eqn.invars), eqn dce_rules[core.free_ref_p] = _free_ref_dce_rule def has_effects(eqn: JaxprEqn, live_ins: Sequence[bool] | None = None) -> bool: """Whether ``eqn`` has effects that prevent it from being DCE'd. If ``live_ins`` is given, Ref effects on inputs that aren't live (i.e. Refs whose contents are never read after ``eqn``) don't count. """ if live_ins is None or all(live_ins): return any(not dceable_effects.contains(e) for e in eqn.effects) dead = {v for v, live in zip(eqn.invars, live_ins) if not live and type(v) is Var} return any(not dceable_effects.contains(e) and not (isinstance(e, RefEffect) and e.input in dead) for e in eqn.effects) @weakref_lru_cache def _dce_jaxpr(jaxpr: Jaxpr, used_outputs: tuple[bool, ...], instantiate: tuple[bool, ...], live_inputs: tuple[bool, ...] ) -> tuple[Jaxpr | None, list[bool]]: env: dict[Var, bool] = {} def read(v: Var) -> bool: return env.get(v, False) def write(x: Atom, b: bool) -> None: if type(x) is Var: env[x] = read(x) or b may_have_dead_refs = (core.internal_mutable_array_effect in jaxpr.effects or not all(live_inputs)) live_at_exit = ({*jaxpr.constvars, *(v for v, l in zip(jaxpr.invars, live_inputs) if l)} if may_have_dead_refs else set()) def live(x: Atom) -> bool: return (type(x) is not Var or read(x) or x in live_at_exit or not isinstance(x.aval, AbstractRef)) new_eqns = [] foreach(write, jaxpr.outvars, used_outputs) for eqn in jaxpr.eqns[::-1]: used_outs = map(read, eqn.outvars) live_ins = (map(live, eqn.invars) if may_have_dead_refs else [True] * len(eqn.invars)) rule = dce_rules.get(eqn.primitive, _default_dce_rule) used_ins, new_eqn = rule(used_outs, live_ins, eqn) if new_eqn is not None: new_eqns.append(new_eqn) foreach(write, eqn.invars, used_ins) used_inputs = map(read, jaxpr.invars) used_inputs = map(op.or_, instantiate, used_inputs) invars = [v for v, b in zip(jaxpr.invars, used_inputs) if b] outvars = [v for v, b in zip(jaxpr.outvars, used_outputs) if b] eqns = new_eqns[::-1] jaxpr_effects = make_jaxpr_effects(jaxpr.constvars, invars, outvars, eqns) if (all(used_inputs) and all(used_outputs) and len(eqns) == len(jaxpr.eqns) and all(new_eqn is eqn for new_eqn, eqn in zip(eqns, jaxpr.eqns)) and jaxpr_effects == jaxpr.effects): return None, used_inputs dbg = core.DebugInfo( jaxpr.debug_info.traced_for, jaxpr.debug_info.func_src_info, jaxpr.debug_info.filter_arg_names(used_inputs), jaxpr.debug_info.filter_result_paths(used_outputs)) new_jaxpr = jaxpr.replace(invars=invars, outvars=outvars, eqns=eqns, effects=jaxpr_effects, debug_info=dbg) config.enable_checks.value and core.check_jaxpr(new_jaxpr) return new_jaxpr, used_inputs DCERule = Callable[[list[bool], list[bool], JaxprEqn], tuple[list[bool], JaxprEqn | None]] def _cached_closed_call_dce(jaxpr_, used_outputs: tuple[bool, ...], live_inputs: tuple[bool, ...] | bool = True, ) -> tuple[Jaxpr, list[bool]]: # dce_jaxpr preserves attached consts (constvars are never pruned). return dce_jaxpr(jaxpr_, used_outputs, live_inputs=live_inputs) def dce_jaxpr_closed_call_rule(used_outputs: list[bool], live_ins: list[bool], eqn: JaxprEqn ) -> tuple[list[bool], JaxprEqn | None]: # TODO(mattjj): de-duplicate with above rule? if not any(used_outputs) and not has_effects(eqn, live_ins): return [False] * len(eqn.invars), None jaxpr_ = eqn.params['call_jaxpr'] closed_jaxpr, used_inputs = _cached_closed_call_dce( jaxpr_, tuple(used_outputs), tuple(live_ins)) if closed_jaxpr is jaxpr_ and all(used_inputs) and all(used_outputs): return used_inputs, eqn new_invars = [v for v, used in zip(eqn.invars, used_inputs) if used] effects = core.eqn_effects(closed_jaxpr, new_invars) new_params = dict(eqn.params, call_jaxpr=closed_jaxpr) new_eqn = new_jaxpr_eqn( new_invars, [v for v, used in zip(eqn.outvars, used_outputs) if used], eqn.primitive, new_params, effects, eqn.source_info, eqn.ctx) return used_inputs, new_eqn def close_jaxpr(jaxpr: Jaxpr) -> Jaxpr: # Now that Jaxpr and ClosedJaxpr are merged, every jaxpr is closed: the # constvars are exactly the inputs with attached const values. return jaxpr def move_binders_to_front(closed_jaxpr: Jaxpr, to_move: Sequence[bool] ) -> Jaxpr: return _move_binders_to_front(closed_jaxpr, tuple(to_move)) @weakref_lru_cache def _move_binders_to_front(jaxpr: Jaxpr, to_move: tuple[bool, ...] ) -> Jaxpr: assert len(jaxpr.in_avals) == len(to_move) new_invars = _move_to_front(jaxpr.invars, to_move) if jaxpr.debug_info.arg_names is None: new_arg_names = None else: new_arg_names = tuple(_move_to_front(jaxpr.debug_info.arg_names, to_move)) dbg = jaxpr.debug_info._replace(arg_names=new_arg_names) return jaxpr.replace(invars=new_invars, debug_info=dbg) def _move_to_front(lst: Sequence, to_move: Sequence[bool]) -> Sequence: return ([elt for elt, move in zip(lst, to_move) if move] + [elt for elt, move in zip(lst, to_move) if not move]) def move_binders_to_back(closed_jaxpr: Jaxpr, to_move: Sequence[bool] ) -> Jaxpr: assert len(to_move) <= len(closed_jaxpr.invars) to_move = [*to_move] + [False] * (len(closed_jaxpr.invars) - len(to_move)) return move_binders_to_front(closed_jaxpr, map(op.not_, to_move)) class DynamicJaxprTracer(Tracer['DynamicJaxprTrace']): __slots__ = ['val', 'parent', '_debug_info'] _trace: DynamicJaxprTrace def __init__(self, trace: DynamicJaxprTrace, aval: core.AbstractValue, val : Atom, line_info: source_info_util.SourceInfo | None = None, parent : TracingEqn | None = None): # TODO(dougalm): Remove aval. It's redundant now that we have val. Tracer.__init__(self, trace, aval) # slightly faster than super() self._line_info = line_info self._debug_info = self._trace.frame.debug_info # for UnexpectedTracerError self.val = val self.parent = parent def _short_repr(self): return f"JitTracer({self.aval})" def full_lower(self): atom = self.val if isinstance(atom, Literal): return atom.val else: maybe_const = self._trace.frame.constvar_to_val.get(atom) if maybe_const is None: return self else: return core.full_lower(maybe_const) def _contents(self): return () def _origin_msg(self): invar_pos, progenitor_eqns = self._trace.frame.find_progenitors(self) dbg = self._debug_info if dbg is None: return "" origin = ("The error occurred while tracing the function " f"{dbg.func_src_info} for {dbg.traced_for}. ") if invar_pos: try: arg_names = [(dbg.arg_names[i] if dbg.arg_names is not None else "unknown") for i in invar_pos] except IndexError: return "" # TODO(mattjj): figure out when not (invar_pos < len(arg_info)) if len(arg_names) == 1: arg_info_str = f"the argument {arg_names[0]}" elif len(arg_names) == 2: arg_info_str = f"the arguments {arg_names[0]} and {arg_names[1]}" else: *rest, last = arg_names arg_info_str = f"the arguments {', '.join(rest)}, and {last}" origin += ("This concrete value was not available in Python because it " f"depends on the value{'s' if len(invar_pos) > 1 else ''} " f"of {arg_info_str}.") elif progenitor_eqns: msts = [" operation " f"{core.pp_eqn(eqn, core.JaxprPpContext(), core.JaxprPpSettings(print_shapes=True))}\n" f" from line {source_info_util.summarize(eqn.source_info)}" for eqn in progenitor_eqns[:5]] # show at most 5 origin += ("This value became a tracer due to JAX operations on these lines:" "\n\n" + "\n\n".join(msts)) if len(progenitor_eqns) > 5: origin += "\n\n(Additional originating lines are not shown.)" return "\n" + origin def get_const(self): return self._trace.get_const(self) def get_referent(self): frame = self._trace.frame atom = self.val val = frame.constvar_to_val.get(atom) if isinstance(atom, Var) else None return self if val is None else get_referent(val) core.pytype_aval_mappings[DynamicJaxprTracer] = lambda x: x.aval def make_jaxpr_effects(constvars, invars, outvars, eqns) -> effects.Effects: jaxpr_effects = set() input_vars = {*constvars, *invars} mut_arrays = set() for eqn in eqns: if eqn.primitive.ref_allocating: outvar, = eqn.outvars mut_arrays.add(outvar) for eff in eqn.effects: if isinstance(eff, effects.JaxprInputEffect): if eff.input in mut_arrays: continue if eff.input not in input_vars: # TODO(mattjj): ask for forgiveness dbg = type('Fake', (), {'resolve_result_paths': lambda self_: self_, 'assert_arg_names': lambda _, __: None, 'assert_result_paths': lambda _, __: None, })() raise ValueError( f"`JaxprInputEffect` {eff} does not have " f"a corresponding jaxpr input." f"\n Equation: {eqn}\n" f"\n Effects: {eqn.effects}\n" "\n Jaxpr: " f"{core.Jaxpr(constvars, invars, outvars, eqns, set(), dbg)}") jaxpr_effects.add(eff) return jaxpr_effects class JaxprStackFrame: __slots__ = ( 'gensym', 'constid_to_tracer', 'constvar_to_val', 'tracing_eqns', 'invars', 'effects', 'debug_info', 'is_high', 'auto_dce') gensym: Callable[[AbstractValue], Var] constid_to_tracer: WeakValueDictionary[ConstId, DynamicJaxprTracer] constvar_to_val: dict[Var, Any] tracing_eqns: list[ReferenceType[TracingEqn] | TracingEqn] invars: list[Var] effects: core.Effects debug_info: core.DebugInfo | None is_high: bool auto_dce: bool def __init__(self, debug_info: core.DebugInfo | None, auto_dce: bool): self.gensym = core.gensym() self.constid_to_tracer = WeakValueDictionary() self.constvar_to_val = {} self.tracing_eqns = [] # cleared when we pop frame from main self.invars = [] self.effects = set() self.debug_info = debug_info self.is_high = False self.auto_dce = auto_dce def add_eqn(self, eqn: TracingEqn): assert isinstance(eqn, TracingEqn) r = eqn if (eqn.effects or not self.auto_dce) else ref(eqn) self.tracing_eqns.append(r) def get_eqns(self): eqns = [] for tracing_eqn in self.tracing_eqns: e = tracing_eqn() if isinstance(tracing_eqn, ReferenceType) else tracing_eqn if e is None: continue if isinstance(e, TracingEqn): e = JaxprEqn( [t.val for t in e.in_tracers], e.outvars, e.primitive, e.params, e.effects, e.source_info, e.ctx) eqns.append(e) return eqns def to_jaxpr( self, trace: DynamicJaxprTrace, out_tracers: Sequence[Tracer], debug_info: core.DebugInfo, source_info: SourceInfo, ) -> tuple[Jaxpr, list[Any]]: # We are careful to snapshot constvar_to_val before we call get_eqns(), # to avoid the following scenario: # * we call get_eqns(), snapshotting the equations. # * garbage collection runs, deleting some TracingEqns, which may # transitively free constant Vars. # * we now have equations with dangling Var references. # If we snapshot the Vars first we won't have this problem. constvars, constvals = unzip2(self.constvar_to_val.copy().items()) eqns = self.get_eqns() outvars = [t.val for t in out_tracers] effs = make_jaxpr_effects(constvars, self.invars, outvars, eqns) all_vars = it.chain(constvars, self.invars, outvars) is_high = self.is_high or any(v.aval.is_high for v in all_vars) jaxpr = Jaxpr(constvars, self.invars, outvars, eqns, effs, debug_info, is_high) return jaxpr, list(constvals) def newvar(self, aval): return self.gensym(aval) def find_progenitors(self, tracer): eqns = self.get_eqns() var = tracer.val if not var or isinstance(var, Literal): return None, None active_vars = {var} for eqn in eqns[::-1]: produced = set(eqn.outvars) & active_vars if produced: active_vars.difference_update(produced) active_vars.update({v for v in eqn.invars if type(v) is Var}) invar_positions = [i for i, v in enumerate(self.invars) if v in active_vars] constvars = active_vars & set(self.constvar_to_val.copy()) const_eqns = [eqn for eqn in eqns if any( v in constvars if type(v) is Var else type(v) is Literal for v in eqn.invars)] return invar_positions, const_eqns ConstFoldRule = Callable[ [list[Any | None], Any, list[AbstractValue]], list[Any | None] | None, ] const_fold_rules: dict[Primitive, ConstFoldRule] = {} ForwardingRule = Callable[ [JaxprEqn], tuple[list[int | None], JaxprEqn | None] ] forwarding_rules: dict[Primitive, ForwardingRule] = {} @multi_weakref_lru_cache def _cached_abstract_eval(primitive: core.Primitive, *avals, **params): return primitive.abstract_eval(*avals, **params) def _verify_params_are_hashable( primitive: core.Primitive, params: dict[str, Any]) -> None: for k, v in params.items(): try: hash(v) except TypeError as e: raise TypeError( "As of JAX v0.7, parameters to jaxpr equations must have __hash__ and " f"__eq__ methods. In a call to primitive {primitive}, the value of " f"parameter {k} was not hashable: {v}") from e # We use TracingEqn instead JaxprEqn during tracing to allow automatic # on-the-fly DCE based on Python refcounting. DynamicJaxprTracers point to # TracingEqns which point to DynamicJaxprTracers and unreachable constants can # be freed. @dataclass(slots=True, weakref_slot=True) class TracingEqn: in_tracers: list[DynamicJaxprTracer] outvars: list[Var] primitive: Primitive params: dict[str, Any] effects: core.Effects source_info: source_info_util.SourceInfo ctx: JaxprEqnContext def __init__(self, in_tracers, outvars, primitive, params, effects, source_info, ctx): self.in_tracers = in_tracers self.outvars = outvars self.primitive = primitive self.params = params self.effects = effects self.source_info = source_info self.ctx = ctx # Allow TracingEqn to duck-type JaxpeEqn because some of the forwarding # rules need to work with both. TODO(dougalm): remove this once we fix # forwarding. @property def invars(self): return self.in_tracers class DynamicJaxprTrace(core.Trace): __slots__ = ("frame", "tag", "parent_trace") # Note that tag is only used when DynamicJaxprTrace is associated with a LinearizeTrace; # otherwise it will be undefined. tag: core.TraceTag frame: JaxprStackFrame parent_trace: core.Trace | None requires_low: bool def __init__(self, debug_info: core.DebugInfo | None, parent_trace: core.Trace | None = None, lower: bool = False, auto_dce: bool =False): super().__init__() self.requires_low = lower self.frame = JaxprStackFrame(debug_info, auto_dce) self.parent_trace = parent_trace def invalidate(self): # TODO(mattjj): exposed existing tracer leaks; fix them and re-enable! # super().invalidate() # avoid cyclic refs self.frame.tracing_eqns = [] # thunk -> eqn -> in_tracers -> trace -> # -> frame -> tracing_eqns -> thunk # TODO(dougalm): we might be able to remove these given refcounting dce self.frame.constid_to_tracer = {} # pyrefly: ignore[bad-assignment] self.frame.constvar_to_val = {} def to_jaxpr_tracer(self, x, source_info: SourceInfo): if isinstance(x, DynamicJaxprTracer) and x._trace is self: return x else: m = getattr(x, "dimension_as_value", None) if m is not None: with core.set_current_trace(self): x = m() return self.to_jaxpr_tracer(x, source_info) else: return self.new_const(x, source_info) def var_to_tracer(self, var, source_info, parent=None): return DynamicJaxprTracer(self, var.aval, var, source_info, parent) def new_arg(self, aval, source_info: SourceInfo): var = self.frame.newvar(aval) tracer = DynamicJaxprTracer(self, aval, var, source_info) self.frame.invars.append(var) return tracer def make_eqn(self, in_tracers, out_avals, primitive, params, effects, source_info=None, ctx = None): source_info = source_info or source_info_util.new_source_info() ctx = ctx or core.current_jaxpr_eqn_context() outvars = map(self.frame.newvar, out_avals) if effects: effects = core.resolve_input_effects(effects, [t.val for t in in_tracers]) if config.enable_checks.value: assert all(isinstance(x, DynamicJaxprTracer) for x in in_tracers) assert all(isinstance(v, Var) for v in outvars) eqn = TracingEqn(in_tracers, outvars, primitive, params, effects, source_info, ctx) out_tracers = [self.var_to_tracer(v, source_info, eqn) for v in outvars] return eqn, out_tracers def emit_eqn(self, in_tracers, out_avals, primitive, params, effects, source_info=None, ctx=None): eqn, out_tracers = self.make_eqn(in_tracers, out_avals, primitive, params, effects, source_info, ctx) self.frame.add_eqn(eqn) return out_tracers def new_const(self, c, source_info: SourceInfo, aval: AbstractValue | None = None): # TODO(mattjj): for ints, or hashable consts, don't rely on id tracer = self.frame.constid_to_tracer.get(id(c)) if tracer is None: if aval is None: aval = typeof(c) tracer = self._new_const(aval, c, source_info) return tracer pure = lift = new_const def _new_const(self, aval, c, source_info: SourceInfo) -> DynamicJaxprTracer: id_c = id(c) assert type(c) not in (int, float, complex, np.generic, np.ndarray), ( f"non-canonical constant of type {type(c).__name__}: {c!r}") if core.is_literalable(c, self.frame.auto_dce): val = Literal(c, aval) return DynamicJaxprTracer(self, aval, val, source_info) else: var = self.frame.newvar(aval) tracer = DynamicJaxprTracer(self, aval, var, source_info) self.frame.constid_to_tracer[id_c] = tracer self.frame.constvar_to_val[var] = c finalize(tracer, self.finalize_const, var, id_c) return tracer def finalize_const(self, var, constid): self.frame.constvar_to_val.pop(var, None) def get_const(self, tracer) -> Any: atom = tracer.val if isinstance(atom, Literal): return atom.val else: return self.frame.constvar_to_val.get(atom) def stage_value(self, val): if config.eager_constant_folding.value and not isinstance(val, Tracer): return core.eval_trace.stage_value(val) source_info = source_info_util.current() return self.to_jaxpr_tracer(val, source_info=source_info) def process_primitive(self, primitive, tracers, params, /): self.frame.is_high |= primitive.is_high(*map(typeof, tracers), **params) if config.eager_constant_folding.value and not any(isinstance(x, Tracer) for x in tracers): avals = tuple(core.typeof(x) for x in tracers) return primitive.bind_with_trace(core.eval_trace, tracers, avals, params) source_info = source_info_util.current() to_jaxpr_tracer = partial(self.to_jaxpr_tracer, source_info=source_info) jaxpr_tracers = map(to_jaxpr_tracer, tracers) if primitive in custom_staging_rules: return custom_staging_rules[primitive](self, source_info, *jaxpr_tracers, **params) return self.default_process_primitive( primitive, jaxpr_tracers, params, source_info) def default_process_primitive(self, primitive, tracers, params, source_info=None): avals = [t.aval for t in tracers] # TODO(mattjj): make custom_lin have hashable params. # TODO(dougalm): add an attribute to primitives to mark primitives with # effectful abstract_eval rules. if (primitive.ref_allocating or primitive.name in ("custom_lin", "call_hi_primitive_linearized", "call_hi_primitive")): out_avals, effs = primitive.abstract_eval(*avals, **params) else: try: out_avals, effs = _cached_abstract_eval(primitive, *avals, **params) except Exception: # TODO(phawkins): remove this 3 months after the release of JAX v0.7. _verify_params_are_hashable(primitive, params) raise if isinstance(out_avals, (tuple, list)) != primitive.multiple_results: raise ValueError(f"{primitive}.abstract_eval() method should return " f"a tuple or a list iff {primitive}.multiple_results.") out_avals = [out_avals] if not primitive.multiple_results else out_avals source_info = source_info or source_info_util.current() maybe_consts_out = try_constant_folding(primitive, tracers, params, out_avals) if maybe_consts_out is not None: eqn = None out_tracers = [self.new_const(c, source_info=source_info, aval=aval) for c, aval in zip(maybe_consts_out, out_avals)] else: eqn, out_tracers = self.make_eqn(tracers, out_avals, primitive, params, effs, source_info=source_info) # Input-to-output tracer forwarding no_input_effects = not any(isinstance(e, effects.JaxprInputEffect) for e in effs) if eqn is not None and no_input_effects and primitive in forwarding_rules: in_fwd, eqn = forwarding_rules[primitive](eqn) for out_idx, in_idx in enumerate(in_fwd): if in_idx is not None: out_tracers[out_idx] = tracers[in_idx] if eqn is not None: self.frame.add_eqn(eqn) # pyrefly: ignore[bad-argument-type] return out_tracers if primitive.multiple_results else out_tracers.pop() def process_custom_jvp_call(self, prim, fun: lu.WrappedFun, jvp: lu.WrappedFun, tracers, /, *, symbolic_zeros: bool): if self.requires_low: with core.set_current_trace(self): return fun.call_wrapped(*tracers) if config.eager_constant_folding.value and not any(isinstance(x, Tracer) for x in tracers): avals = tuple(core.typeof(x) for x in tracers) return prim.bind_with_trace( core.eval_trace, tracers, avals, dict(subfuns=(fun, jvp), symbolic_zeros=symbolic_zeros)) source_info = source_info_util.current() to_jaxpr_tracer = partial(self.to_jaxpr_tracer, source_info=source_info) tracers = map(to_jaxpr_tracer, tracers) in_avals = [t.aval for t in tracers] in_tangent_avals = [t.to_tangent_aval() for t in in_avals] fun_jaxpr, out_avals, consts = trace_to_jaxpr_dynamic(fun, in_avals, lower=self.requires_low) self.frame.is_high |= fun_jaxpr.is_high closed_fun_jaxpr = convert_constvars_jaxpr(fun_jaxpr) @partial(lu.wrap_init, debug_info=jvp.debug_info) @_memoize def jvp_jaxpr_thunk(*in_zeros): for store in jvp.stores: store and store.reset() nz_tangent_avals, zero_avals = partition_list(in_zeros, in_tangent_avals) jvp_, out_zeros = _jvp_jaxpr_zeros(jvp, in_zeros, tuple(zero_avals)) in_avals_ = (*in_avals, *nz_tangent_avals) jaxpr, _, out_consts = trace_to_jaxpr_dynamic(jvp_.with_unknown_names(), in_avals_, lower=self.requires_low) return jaxpr, out_consts, out_zeros() const_tracers = map(to_jaxpr_tracer, consts) return self.emit_eqn( [*const_tracers, *tracers], out_avals, prim, dict(call_jaxpr=closed_fun_jaxpr, jvp_jaxpr_fun=jvp_jaxpr_thunk, num_consts=len(consts), symbolic_zeros=symbolic_zeros), core.positional_effects(closed_fun_jaxpr), source_info=source_info) def process_custom_vjp_call(self, prim: core.Primitive, fun: lu.WrappedFun, fwd: lu.WrappedFun, bwd: lu.WrappedFun, tracers, /, *, out_trees: Callable[[], tuple[PyTreeDef, PyTreeDef, list[int | None]]], symbolic_zeros: bool): if self.requires_low: with core.set_current_trace(self): return fun.call_wrapped(*tracers) if config.eager_constant_folding.value and not any(isinstance(x, Tracer) for x in tracers): avals = tuple(core.typeof(x) for x in tracers) return prim.bind_with_trace( core.eval_trace, tuple(tracers), avals, dict(subfuns=(fun, fwd, bwd), out_trees=out_trees, symbolic_zeros=symbolic_zeros)) source_info = source_info_util.current() to_jaxpr_tracer = partial(self.to_jaxpr_tracer, source_info=source_info) tracers = map(to_jaxpr_tracer, tracers) in_avals = [t.aval for t in tracers] fun_jaxpr, out_avals, consts = trace_to_jaxpr_dynamic(fun.with_unknown_names(), in_avals, lower=self.requires_low) if config.mutable_array_checks.value: api_util._check_no_aliased_closed_over_refs(fun.debug_info, consts, tracers) self.frame.is_high |= fun_jaxpr.is_high num_consts = len(consts) closed_fun_jaxpr = convert_constvars_jaxpr(fun_jaxpr) @partial(lu.wrap_init, debug_info=fwd.debug_info) @_memoize def fwd_jaxpr_from_zeros(*zeros): for store in fwd.stores: store and store.reset() fwd_ = _interleave_fun(fwd.with_unknown_names(), zeros) jaxpr, _, consts = trace_to_jaxpr_dynamic(fwd_, in_avals, lower=self.requires_low) return jaxpr, consts def out_trees_(): out_tree, res_tree, input_fwds = out_trees() input_fwds = [f if f is None else f + num_consts for f in input_fwds] return out_tree, res_tree, input_fwds const_tracers = map(to_jaxpr_tracer, consts) return self.emit_eqn( [*const_tracers, *tracers], out_avals, prim, dict(call_jaxpr=closed_fun_jaxpr, fwd_jaxpr_thunk=fwd_jaxpr_from_zeros, num_consts=num_consts, bwd=bwd, out_trees=out_trees_, symbolic_zeros=symbolic_zeros), core.positional_effects(closed_fun_jaxpr), source_info=source_info) def to_jaxpr(self, out_tracers: Sequence[Tracer], debug_info: core.DebugInfo, source_info: SourceInfo): return self.frame.to_jaxpr(self, out_tracers, debug_info, source_info) custom_staging_rules: dict[Primitive, Callable] = {} @lu.transformation2 def _interleave_fun(f, every_others, *args, **kwargs): args_ = [x for pair in zip(args, every_others) for x in pair] return f(*args_, **kwargs) # TODO: consider renaming to "lazy_thunk" def _memoize(fn): cells = {} sentinel = object() def memoized(*args): out = cells.get(args, sentinel) if out is sentinel: with core.set_current_trace(None): out = cells[args] = fn(*args) return out return memoized @lu.transformation_with_aux2 def _jvp_jaxpr_zeros(f, store, in_zeros, zero_avals, *primal_tangent_avals): in_primals, nz_in_tangents = split_list(primal_tangent_avals, [len(in_zeros)]) symbolic_zeros = map(ad_util.SymbolicZero, zero_avals) tangents = merge_lists(in_zeros, nz_in_tangents, symbolic_zeros) out = f(*in_primals, *tangents) n, ragged = divmod(len(out), 2) assert not ragged out_primals, out_tangents = out[:n], out[n:] out_zeros = [type(t) is ad_util.SymbolicZero for t in out_tangents] out_nz_tangents, _ = partition_list(out_zeros, out_tangents) store.store(out_zeros) return [*out_primals, *out_nz_tangents] callsites_with_tracing_cache_miss: set[str] = set() def explain(keys, fun, in_avals, debug_info, *context, **_): func_filename = debug_info.func_filename if func_filename and not source_info_util.is_user_filename(func_filename): return msg: list[str] = [] p = msg.append callsite = source_info_util.summarize(source_info_util.current()) p(f"TRACING CACHE MISS at {callsite}:") src_info = "" if func_filename: src_info += f" defined at {func_filename}" if func_lineno := debug_info.func_lineno: src_info += f":{func_lineno}" func_name = debug_info.func_name # have we seen this function before at all? keys = [key for fun_ref, *key in keys if fun_ref() is fun] if not keys: p(f" never seen function:\n {func_name} id={id(fun)}{src_info}") if callsite in callsites_with_tracing_cache_miss: p(" but seen another function defined on the same line; maybe the function is\n" " being re-defined repeatedly, preventing caching?") else: callsites_with_tracing_cache_miss.add(callsite) return logger.log(logging.WARNING, "\n".join(msg)) p(f" for {func_name}{src_info}") key = (config.trace_context(), (in_avals, debug_info, *context), {}) min_diff = min(diff_tracing_cache_keys(key, k) for k in keys)[-1] p(' all previously seen cache keys differ. For the closest previous key:') p(' ' + min_diff) return logger.log(logging.WARNING, "\n".join(msg)) def diff_tracing_cache_keys(new_key, old_key) -> tuple[int, int, str]: """Explain the diff between two tracing cache keys. Returns: A tuple of (severity, num_diffs, explanation) for the diff between the two keys. Severity is an int, where lower is better. """ new_ctx, (new_tree, new_dbg, *_), () = new_key old_ctx, (old_tree, old_dbg, *_), () = old_key return (diff_tracing_ctx(new_ctx, old_ctx) or diff_trees(new_tree.tree, old_tree.tree) or diff_debug(new_dbg, old_dbg) or diff_types(new_dbg, new_tree.vals, old_tree.vals) or (4, 0, 'cache miss explanation unavailable')) def diff_tracing_ctx(new_ctx, old_ctx) -> tuple[int, int, str] | None: if new_ctx == old_ctx: return None diffs: list[str] = [] msg = "Tracing context doesn't match, e.g. due to config or context manager." if len(new_ctx) != len(old_ctx): num_diff = abs(len(new_ctx) - len(old_ctx)) diffs.append(" * number of tracing context values differs: " f"now {len(new_ctx)} and before {len(old_ctx)}") return 0, num_diff, msg + "\n" + "\n".join(diffs) num_diff = sum(map(op.ne, new_ctx, old_ctx)) context_names = config.trace_context_names() if len(context_names) != len(new_ctx): diffs.append(" * number of tracing context names differs: " f"context_names {len(context_names)} vs " f"context length {len(new_ctx)}") return 0, num_diff, msg + "\n" + "\n".join(diffs) for name, new, old in zip(context_names, new_ctx, old_ctx): if new != old: diffs.append(f" * {name} differs: now {new} and before {old}") return 0, num_diff, msg + "\n" + "\n".join(diffs) def diff_trees(new_tree, old_tree) -> tuple[int, int, str] | None: errs = tree_util.equality_errors_pytreedef(new_tree, old_tree) tree_diffs = [] for path, thing1, thing2, explanation in errs: tree_diffs.append( f" * at input path {tree_util.keystr(tuple(path))}, now {thing1} and " f"before {thing2}, so {explanation}") msg = 'different input pytree:\n' + '\n'.join(tree_diffs) if tree_diffs: return 1, len(tree_diffs), msg def diff_debug(new_dbg, old_dbg) -> tuple[int, int, str] | None: msg = "Debug info doesn't match." num_diff = sum(map(op.ne, new_dbg, old_dbg)) if num_diff: return 2, num_diff, msg def diff_types(dbg, new_leaves, old_leaves) -> tuple[int, int, str] | None: if new_leaves == old_leaves: return diffs = [] add_weak_type_hint = False for name, new_ty, old_ty in zip(dbg.arg_names, new_leaves, old_leaves): if new_ty != old_ty: new_str, old_str = new_ty.str_short(True), old_ty.str_short(True) if type(new_ty) is type(old_ty) is core.ShapedArray: if new_ty.sharding != old_ty.sharding: new_str, old_str = new_ty.str_short(True, True), old_ty.str_short(True, True) if new_ty.weak_type != old_ty.weak_type: add_weak_type_hint = True new_str += f'{{weak_type={new_ty.weak_type}}}' old_str += f'{{weak_type={old_ty.weak_type}}}' diffs.append(f" * at {name}, now {new_str} and before {old_str}") msg = 'different input types:\n' + '\n'.join(diffs) if add_weak_type_hint: msg += 'https://docs.jax.dev/en/latest/101/type_promotion.html#weak-types' if diffs: return 3, len(diffs), msg def _lower_debug_info(hi_jaxpr): debug_info = hi_jaxpr.debug_info if debug_info.arg_names is not None: lo_arg_names = tuple( name for aval, name in zip(hi_jaxpr.in_avals, debug_info.arg_names) for _ in aval.lo_ty()) debug_info = debug_info._replace(arg_names=lo_arg_names) if debug_info.result_paths is not None: lo_result_paths = tuple( path for aval, path in zip(hi_jaxpr.out_avals, debug_info.result_paths) for _ in aval.lo_ty()) debug_info = debug_info._replace(result_paths=lo_result_paths) return debug_info def trace_to_jaxpr_nocache( fun: Callable, in_avals: ft.FlatTree, # (args, kwargs) pair debug_info: core.DebugInfo, # TODO: let's just make a `trace_to_jaxpr_ft` function for this fun_takes_flat_tree_arg=False, fun_returns_flat_tree=False, requires_low=False, ) -> tuple[Jaxpr, ft.FlatTree]: if config.no_tracing.value: raise RuntimeError(f"re-tracing function {fun} for " "`jit`, but 'no_tracing' is set") test_event("trace_to_jaxpr") config.enable_checks.value and debug_info.assert_arg_names(len(in_avals)) parent_trace = core.trace_ctx.trace trace = DynamicJaxprTrace(debug_info, parent_trace=parent_trace, lower=requires_low) # Name stack and the traceback scope are reset because the metadata on jaxpr # equations should be rooted at the enclosing jaxpr and not contain any # context from the callsite. Otherwise metadata from one caller would bleed # into metadata from a different caller if we, e.g., inline. with (core.ensure_no_leaks(trace), source_info_util.reset_name_stack(), TracebackScope()): source_info = source_info_util.current() if requires_low: if debug_info.arg_names is not None: debug_info = debug_info._replace(arg_names=tuple( name for aval, name in zip(in_avals, debug_info.arg_names) for _ in aval.lo_ty())) def new_arg(aval): lo_tracers = [trace.new_arg(lo_aval, source_info=source_info) for lo_aval in aval.lo_ty()] # noqa: F821 return aval.raise_val(*lo_tracers) in_tracers = in_avals.map(new_arg) else: in_tracers = in_avals.map(partial(trace.new_arg, source_info=source_info)) with core.set_current_trace(trace): if fun_takes_flat_tree_arg: args_ft, kwargs_ft = in_tracers.unpack() assert kwargs_ft.unflatten() == {} # TODO: handle kwargs kwargs = {} args = args_ft.unpack() del args_ft else: args, kwargs = in_tracers.unflatten() ans_pytree = fun(*args, **kwargs) if fun_returns_flat_tree: # TODO(dougalm): make result paths optional ans = ans_pytree debug_info = debug_info.set_result_paths([''] * len(ans)) else: debug_info = debug_info.set_result_paths(ans_pytree) ans = ft.flatten(ans_pytree) del ans_pytree, args, kwargs _check_returned_jaxtypes(debug_info, list(ans)) ans = ans.map(dtypes.canonicalize_value) out_avals = ans.map(typeof) if requires_low: flat_out_tracers = [trace.to_jaxpr_tracer(x, source_info=source_info) for aval, hi_val in zip(out_avals, ans) for x in aval.lower_val(hi_val)] debug_info = debug_info.resolve_result_paths() if debug_info.result_paths is not None: debug_info = debug_info._replace(result_paths=tuple( path for aval, path in zip(out_avals, debug_info.result_paths) # pyrefly: ignore[bad-argument-type] for _ in aval.lo_ty())) else: flat_out_tracers = [trace.to_jaxpr_tracer(x, source_info=source_info) for x in ans] _check_no_returned_refs(debug_info, list(flat_out_tracers)) jaxpr, consts = trace.frame.to_jaxpr(trace, list(flat_out_tracers), debug_info, source_info) del trace, fun, in_tracers, flat_out_tracers, ans config.enable_checks.value and core.check_jaxpr(jaxpr) return jaxpr.with_consts(consts), out_avals trace_to_jaxpr = weakref_lru_cache(maxsize=None, explain=explain)( trace_to_jaxpr_nocache) # TODO(dougalm): remove in favor of `trace_to_jaxpr` @profiler.annotate_function def trace_to_jaxpr_dynamic( fun: lu.WrappedFun, in_avals: Sequence[AbstractValue], *, keep_inputs: list[bool] | None = None, lower: bool = False, auto_dce: bool = False) -> tuple[Jaxpr, list[AbstractValue], list[Any]]: config.enable_checks.value and fun.debug_info.assert_arg_names(len(in_avals)) keep_inputs = [True] * len(in_avals) if keep_inputs is None else keep_inputs parent_trace = core.trace_ctx.trace trace = DynamicJaxprTrace(fun.debug_info, parent_trace=parent_trace, lower=lower, auto_dce=auto_dce) # Name stack and the traceback scope are reset because the metadata on jaxpr # equations should be rooted at the enclosing jaxpr and not contain any # context from the callsite. Otherwise metadata from one caller would bleed # into metadata from a different caller if we, e.g., inline. with (core.ensure_no_leaks(trace), source_info_util.reset_name_stack(), TracebackScope()): source_info = source_info_util.current() in_tracers = map(partial(trace.new_arg, source_info=source_info), in_avals) in_tracers = [t for t, keep in zip(in_tracers, keep_inputs) if keep] with core.set_current_trace(trace): ans = fun.call_wrapped(*in_tracers) _check_returned_jaxtypes(fun.debug_info, ans) ans = map(dtypes.canonicalize_value, ans) out_tracers = map(partial(trace.to_jaxpr_tracer, source_info=source_info), ans) _check_no_returned_refs(fun.debug_info, out_tracers) jaxpr, consts = trace.frame.to_jaxpr(trace, out_tracers, fun.debug_info, source_info) del trace, fun, in_tracers, out_tracers, ans config.enable_checks.value and core.check_jaxpr(jaxpr) return jaxpr, [v.aval for v in jaxpr.outvars], consts def _check_returned_jaxtypes(dbg, out_tracers): for i, x in enumerate(out_tracers): try: typeof(x) except TypeError: if (dbg and len(paths := dbg.resolve_result_paths().result_paths) > i and (p := paths[i].removeprefix('result'))): extra = f' at output component {p}' else: extra = '' raise TypeError( f"function {dbg.func_src_info} traced for {dbg.traced_for} returned a " f"value of type {type(x)}{extra}, which is not a valid JAX type") from None def _check_no_returned_refs( dbg: core.DebugInfo, out_tracers: Sequence[DynamicJaxprTracer] ) -> None: if not config.mutable_array_checks.value: return for i, t in enumerate(out_tracers): a = t.aval if isinstance(a, AbstractRef): result_paths = dbg.resolve_result_paths().safe_result_paths(len(out_tracers)) if list(result_paths) == ["result"]: result_paths = [""] # TODO(mattjj): fix in callee loc = result_paths[i] and f' at output tree path {result_paths[i]}' frame = t._trace.frame v = t.val eqns = frame.get_eqns() # TODO(dougalm): something more efficient eqn = next((e for e in eqns if v in e.outvars), None) if eqn: assert eqn.primitive.ref_allocating origin_info = ('\n\nThe returned mutable array was created on line ' f'{source_info_util.summarize(eqn.source_info)}.') elif v in frame.invars: assert isinstance(v, Var) arg_name = dbg.safe_arg_names(len(frame.invars))[frame.invars.index(v)] origin_info = ('\n\nThe returned mutable array was passed in as the ' f'argument {arg_name}.') else: origin_info = ('\n\nThe returned mutable array was closed over by the ' 'function, rather than passed in as an argument or ' 'created inside it.') raise ValueError( f"function {dbg.func_src_info} traced for {dbg.traced_for} returned " f"a mutable array reference of type {a.str_short()}{loc}, but " f"mutable array references cannot be returned.{origin_info}") class TracerAsName: ref: Any def __init__(self, tracer): self.ref = core.get_referent(tracer) def __eq__(self, other): return isinstance(other, TracerAsName) and self.ref is other.ref def __hash__(self): return id(self.ref) Const = Any Val = Any def inline_jaxpr_into_trace( trace: DynamicJaxprTrace, src: SourceInfo, jaxpr: Jaxpr, consts: Sequence[Any], *arg_tracers: DynamicJaxprTracer) -> list[Any]: # This function is conceptually the same thing as just calling eval_jaxpr, const_tracers = map(partial(trace.new_const, source_info=src), consts) env: dict[Var, DynamicJaxprTracer] = dict( zip([*jaxpr.constvars, *jaxpr.invars], [*const_tracers, *arg_tracers])) def inline_atom(src_, x): if isinstance(x, Literal): return DynamicJaxprTracer(trace, x.aval, x, src_) else: return env[x] for eqn in jaxpr.eqns: src_ = (src if not eqn.source_info.name_stack else src.replace(name_stack=src.name_stack + eqn.source_info.name_stack)) in_tracers = map(partial(inline_atom, src_), eqn.invars) out_avals = [v.aval for v in eqn.outvars] maybe_consts = try_constant_folding(eqn.primitive, in_tracers, eqn.params, out_avals) if maybe_consts is not None: out_tracers = [trace.new_const(c, source_info=src_, aval=aval) for c, aval in zip(maybe_consts, out_avals)] else: effs = {e.replace(env[e.input].val) if isinstance(e, effects.JaxprInputEffect) else e for e in eqn.effects} if eqn.effects else eqn.effects out_tracers = trace.emit_eqn(in_tracers, out_avals, eqn.primitive, eqn.params, effs, src_, eqn.ctx) foreach(env.setdefault, eqn.outvars, out_tracers) return map(partial(inline_atom, src), jaxpr.outvars) def try_constant_folding(primitive, tracers, params, out_avals): if primitive in const_fold_rules: consts_in = [t.get_const() for t in tracers] if any(c is not None for c in consts_in): return const_fold_rules[primitive](consts_in, params, out_avals) return None @weakref_lru_cache def lower_jaxpr2(hi_jaxpr) -> Jaxpr: in_avals = ft.flatten(([a.lo_ty() for a in hi_jaxpr.in_avals], {})) lo_jaxpr, _ = lower_jaxpr(hi_jaxpr, in_avals) return lo_jaxpr @weakref_lru_cache def lower_jaxpr(hi_jaxpr: Jaxpr, lo_avals) -> tuple[Jaxpr, ft.FlatTree]: # TODO(yashkatariya): Remove this once we have XLA DCE opaque opt-barrier. hi_jaxpr, _ = dce_jaxpr(hi_jaxpr, True, instantiate=True) env: dict[Var, DynamicJaxprTracer | HTLV] = {} # type:ignore parent_trace = core.trace_ctx.trace trace = DynamicJaxprTrace(hi_jaxpr.debug_info.with_unknown_names(), parent_trace=parent_trace, lower=True) def read(src, x): if isinstance(x, Literal): return x.val elif (htlv := env.get(x)) is not None: # noqa return htlv else: assert not x.aval.is_high return trace.var_to_tracer(x, src) # noqa def write(v, x): if v.aval.is_high: env[v] = x # noqa: F821 elif isinstance(x, DynamicJaxprTracer) and x._trace is trace: # noqa: F821 env[v] = x # noqa: F821 else: env[v] = lift_lo_const(v, x) # noqa: F821 def lift_lo_const(v, c) -> DynamicJaxprTracer: tracer = DynamicJaxprTracer(trace, v.aval, v, src) # noqa: F821 trace.frame.constid_to_tracer[id(c)] = tracer # noqa: F821 trace.frame.constvar_to_val[v] = c # noqa: F821 return tracer with (core.ensure_no_leaks(trace), source_info_util.reset_name_stack(), TracebackScope()): src = source_info_util.current() invals = outs = maybe_invals = new_invars = xs = t = None lo_avals_lol, () = lo_avals.unflatten() for v, xs in zip(hi_jaxpr.invars, lo_avals_lol): if v.aval.is_high: xs = [trace.new_arg(x, source_info=src) for x in xs] env[v] = v.aval.raise_val(*xs) # type: ignore else: trace.frame.invars.append(v) for v, c in zip(hi_jaxpr.constvars, hi_jaxpr.consts): if v.aval.is_high: env[v] = c # treated as an HTLV else: env[v] = lift_lo_const(v, c) with core.set_current_trace(trace): eqns = trace.frame.tracing_eqns for eqn in hi_jaxpr.eqns: maybe_invals = [env.get(x) if isinstance(x, Var) else None for x in eqn.invars] hi = eqn.primitive.is_high(*[v.aval for v in eqn.invars], **eqn.params) if all(x is None for x in maybe_invals) and not hi: eqns.append(eqn) # type: ignore elif not hi: new_invars = [x if isinstance(x, Literal) else t.val if (t := env.get(x)) is not None else x for x in eqn.invars] eqns.append(eqn.replace(invars=new_invars)) else: invals = map(partial(read, eqn.source_info), eqn.invars) name_stack = source_info_util.current_name_stack() + eqn.source_info.name_stack with (source_info_util.user_context(eqn.source_info.traceback, name_stack=name_stack), eqn.ctx.manager): outs = eqn.primitive.to_lojax(*invals, **eqn.params) foreach(write, eqn.outvars, outs if eqn.primitive.multiple_results else [outs]) tracer = partial(trace.to_jaxpr_tracer, source_info=src) out_tracers = [dtypes.canonicalize_value(read(src, x)) for x in hi_jaxpr.outvars] out_tracers = [v.aval.lower_val2(hi_val).map(tracer) for v, hi_val in zip(hi_jaxpr.outvars, out_tracers)] out_tracers = ft.pack(tuple(out_tracers)) out_avals = out_tracers.map(typeof) dbg = _lower_debug_info(hi_jaxpr) jaxpr, consts = trace.frame.to_jaxpr(trace, list(out_tracers), dbg, src) del trace, env, out_tracers, tracer, read, outs, invals, eqns del maybe_invals, new_invars, xs, t config.enable_checks.value and core.check_jaxpr(jaxpr) assert not any(v.aval.is_high for v in it.chain(jaxpr.constvars, jaxpr.invars)) return jaxpr.with_consts(consts), out_avals @weakref_lru_cache def lower_jaxpr_reference(hi_jaxpr: Jaxpr, lo_avals) -> tuple[Jaxpr, ft.FlatTree]: """Reference implementation of lower_jaxpr, for testing and debugging. Swap it in with `pe.lower_jaxpr = pe.lower_jaxpr_reference`; all call sites look the name up at call time. """ dbg = _lower_debug_info(hi_jaxpr) return trace_to_jaxpr(partial(_lower_traceable, hi_jaxpr), lo_avals, dbg, requires_low=True, fun_returns_flat_tree=True) def _lower_traceable(jaxpr, *lo_args): hi_args = [a.raise_val(*xs) for a, xs in zip(jaxpr.in_avals, lo_args)] hi_outs = core.jaxpr_as_fun(jaxpr)(*hi_args) lo_outs = [a.lower_val2(y) for a, y in zip(jaxpr.out_avals, hi_outs)] return ft.pack(tuple(lo_outs)) # vestigial hijax helpers def raise_lo_outs(hi_avals, lo_outs): lo_outs_ = iter(lo_outs) hi_outs = [t.raise_val(*it.islice(lo_outs_, len(t.lo_ty()))) for t in hi_avals] assert next(lo_outs_, None) is None return hi_outs eval_jaxpr_p = core.eval_jaxpr_p eval_jaxpr_p.skip_canonicalization = True dce_rules[eval_jaxpr_p] = dce_jaxpr_closed_call_rule dce_jaxpr_call_rule = dce_jaxpr_closed_call_rule # alias for downstream users def _eval_jaxpr_partial_eval(prim, trace, *in_tracers, call_jaxpr, **params): in_pvals = [t.pval for t in in_tracers] unknown_ins = [not pv.is_known() for pv in in_pvals] known_jaxpr, unknown_jaxpr, unknown_outs, res_avals, in_fwd_res = \ partial_eval_jaxpr_nounits_fwd(call_jaxpr, unknown_ins, instantiate=False) consts = [pv.get_known() for pv in in_pvals if pv.is_known()] all_known_outs = prim.bind(*consts, call_jaxpr=known_jaxpr, **params) known_outs, res = split_list(all_known_outs, [len(all_known_outs) - len(res_avals)]) res_ = iter(res) res = [next(res_) if f is None else [*call_jaxpr.consts, *consts][f] for f in in_fwd_res] assert next(res_, sentinel := object()) is sentinel res_tracers = map(trace.new_instantiated_const, res) unk_tracers_in = [t for t in in_tracers if not t.pval.is_known()] unk_tracers_out = [JaxprTracer(trace, PartialVal.unknown(aval), None) for aval in unknown_jaxpr.out_avals] eqn = new_eqn_recipe(trace, [*res_tracers, *unk_tracers_in], unk_tracers_out, prim, dict(params, call_jaxpr=unknown_jaxpr), core.positional_effects(unknown_jaxpr), source_info_util.current()) for t in unk_tracers_out: t.recipe = eqn if effects.partial_eval_kept_effects.filter_in(unknown_jaxpr.effects): trace.effect_handles.append(EffectHandle([*unk_tracers_in, *res_tracers], eqn)) # type: ignore return merge_lists(unknown_outs, known_outs, unk_tracers_out) custom_partial_eval_rules[eval_jaxpr_p] = partial(_eval_jaxpr_partial_eval, eval_jaxpr_p) partial_eval_jaxpr_custom_rules[eval_jaxpr_p] = \ partial(closed_call_partial_eval_custom_rule, 'call_jaxpr', lambda _, __, ___, ____, _____, ______, x, y: (x, y)) def _lower_and_eval(prim, jaxpr: Jaxpr, args: Sequence[Any], **params): lo_jaxpr = lower_jaxpr2(jaxpr) lo_args = [lo_val for aval, x in zip(jaxpr.in_avals, args) for lo_val in aval.lower_val(x)] lo_outs = prim.bind(*lo_args, call_jaxpr=lo_jaxpr, **params) lo_outs_ = iter(lo_outs) hi_outs = [t.raise_val(*it.islice(lo_outs_, len(t.lo_ty()))) for t in jaxpr.out_avals] assert next(lo_outs_, None) is None return hi_outs def _eval_jaxpr_to_lojax(prim, *hi_args, call_jaxpr: Jaxpr, **params): return _lower_and_eval(prim, call_jaxpr, hi_args, **params) eval_jaxpr_p.to_lojax = partial(_eval_jaxpr_to_lojax, eval_jaxpr_p)