# 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 Counter, defaultdict, deque, namedtuple from collections.abc import (Callable, Collection, Hashable, Iterable, Iterator, Sequence, MutableSet, MutableMapping) from contextlib import contextmanager from dataclasses import dataclass import enum import functools from functools import partial, total_ordering import gc import inspect import itertools as it import math import operator import re import sys import threading import types from typing import (Any, ClassVar, NamedTuple, final, overload, TYPE_CHECKING, Literal as Literal_) import warnings import weakref import numpy as np from jax._src import dtypes from jax._src import config from jax._src import effects from jax._src.frozen_dict import FrozenDict from jax._src import mesh as mesh_lib from jax._src.mesh import AxisType from jax._src.partition_spec import PartitionSpec as P, UnreducedKind from jax._src.errors import ( ConcretizationTypeError, TracerArrayConversionError, TracerBoolConversionError, TracerIntegerConversionError, UnexpectedTracerError) from jax._src import flattree as ft from jax._src import linear_util as lu from jax._src.tree_util import tree_map from jax._src import source_info_util from jax._src.util import (safe_zip, safe_map, curry, tuple_insert, tuple_delete, cache, HashableWrapper, weakref_lru_cache, partition_list, StrictABCMeta, foreach, weakref_cache_key_types, set_module, weak_value_interner, immutable) import jax._src.pretty_printer as pp from jax._src.named_sharding import NamedSharding, get_replicated_axes from jax._src import named_sharding as ns from jax._src.sharding import Sharding from jax._src.layout import (Format, AutoLayout, AutoLayoutSingleton, get_layout_mode, LayoutMode, layout_tracing_mode) from jax._src.lib import _jax from jax._src import traceback_util from jax._src.typing import Array, ArrayLike, DimSize, Shape from jax._src.xla_metadata_lib import current_xla_metadata, update_metadata traceback_util.register_exclusion(__file__) zip, unsafe_zip = safe_zip, zip map, unsafe_map = safe_map, map config_ext = _jax.config PyTree = Any _TRACER_ERROR_NUM_TRACEBACK_FRAMES = config.int_flag( 'jax_tracer_error_num_traceback_frames', config.int_env('JAX_TRACER_ERROR_NUM_TRACEBACK_FRAMES', 5), help='Set the number of stack frames in JAX tracer error messages.' ) def identity(x): return x # -------------------- jaxprs -------------------- Effect = effects.Effect Effects = effects.Effects EffectTypeSet = effects.EffectTypeSet no_effects: Effects = effects.no_effects DebugInfo = lu.DebugInfo InitialResultPaths = lu.InitialResultPaths initial_result_paths = lu.initial_result_paths class Jaxpr: __slots__ = [ "__weakref__", "_all_invars", "_outvars", "_eqns", "_effects", "_debug_info", "_is_high", "_consts", ] _all_invars: list[Var] _outvars: list[Atom] _eqns: list[JaxprEqn] _effects: Effects _debug_info: DebugInfo _is_high: bool _consts: list[Any] @property def all_invars(self) -> list[Var]: return self._all_invars @property def constvars(self) -> list[Var]: # The constant inputs are exactly those with values attached; an input # without an attached value is a plain invar. return self._all_invars[: len(self._consts)] @property def consts(self) -> list[Any]: return self._consts literals = consts # legacy ClosedJaxpr name for consts @property def num_consts(self) -> int: return len(self._consts) @property def jaxpr(self) -> Jaxpr: # Legacy accessor from the days of ClosedJaxpr, which wrapped a Jaxpr. # TODO(dougalm): remove uses and delete. return self @property def invars(self) -> list[Var]: num_consts = len(self._consts) return self._all_invars[num_consts:] if num_consts else self._all_invars @property def outvars(self) -> list[Atom]: return self._outvars @property def eqns(self) -> list[JaxprEqn]: return self._eqns @property def effects(self) -> Effects: return self._effects @property def debug_info(self) -> DebugInfo: return self._debug_info @property def is_high(self) -> bool: return self._is_high @property def in_avals(self): return [v.aval for v in self.invars] @property def out_avals(self): return [v.aval for v in self.outvars] def __init__( self, constvars: Sequence[Var] | Jaxpr, invars: Sequence[Var] | Sequence[Any] | None = None, outvars: Sequence[Atom] | None = None, eqns: Sequence[JaxprEqn] | None = None, effects: Effects = no_effects, # We want all calls to pass a DebugInfo object, but for backwards # compatibility we have to allow calls when the debug_info # is missing. debug_info: DebugInfo = None, # pyrefly: ignore[bad-function-definition] is_high: bool = False, consts: Sequence[Any] | None = None, ): if isinstance(constvars, Jaxpr): # Legacy ClosedJaxpr(jaxpr, consts) construction: share `jaxpr`'s # structure and attach `consts` as its constant argument values. # TODO(dougalm): migrate callers and remove. jaxpr = constvars assert outvars is None and eqns is None and debug_info is None if consts is None: jaxpr, consts = constvars, invars else: assert invars is None assert consts is not None and len(consts) <= len(jaxpr._all_invars) self._all_invars = jaxpr._all_invars self._outvars = jaxpr._outvars self._eqns = jaxpr._eqns self._effects = jaxpr._effects self._debug_info = _shift_arg_names(jaxpr._debug_info, len(jaxpr._consts) - len(consts)) self._is_high = jaxpr._is_high self._consts = list(consts) return assert invars is not None and outvars is not None and eqns is not None self._consts = [] if consts is None else list(consts) assert ( not self._consts or len(self._consts) == len(constvars) ), "consts, when attached, must pair with constvars" self._all_invars = [*constvars, *invars] self._outvars = list(outvars) self._eqns = list(eqns) self._effects = effects # TODO(https://github.com/jax-ml/jax/issues/26480) debug_info = debug_info or lu._missing_debug_info("core.Jaxpr") debug_info = debug_info.resolve_result_paths() if constvars and not self._consts: # Constvars without attached values are plain leading invars. debug_info = _shift_arg_names(debug_info, len(constvars)) self._debug_info = debug_info config.enable_checks.value and self._debug_info.assert_arg_names(len(self.invars)) config.enable_checks.value and self._debug_info.assert_result_paths(len(outvars)) self._is_high = is_high def __str__(self): return str(self.pretty_print()) __repr__ = __str__ def pretty_print(self, *, source_info=False, print_shapes=True, custom_pp_eqn_rules=True, name_stack=False, print_effects: bool = False, **kwargs): doc = pp_toplevel_jaxpr( self, source_info=source_info, print_shapes=print_shapes, custom_pp_eqn_rules=custom_pp_eqn_rules, name_stack=name_stack, print_effects=print_effects) return doc.format(**kwargs) def _repr_pretty_(self, p, cycle): return p.text(self.pretty_print(use_color=True)) def with_consts(self, consts: Sequence[Any]) -> Jaxpr: """Returns a copy of this jaxpr with `consts` attached as the values of its first `len(consts)` inputs, which thereby become its constvars. Shares all other structure.""" consts = list(consts) assert len(consts) <= len(self._all_invars) new = Jaxpr.__new__(Jaxpr) new._all_invars = self._all_invars new._outvars = self._outvars new._eqns = self._eqns new._effects = self._effects new._debug_info = _shift_arg_names(self._debug_info, len(self._consts) - len(consts)) new._is_high = self._is_high new._consts = consts return new def map_jaxpr(self, f): # Legacy ClosedJaxpr method: apply f to the jaxpr, keeping consts. return Jaxpr(f(self), self.consts) def replace(self, **kwargs): if "jaxpr" in kwargs: # Legacy ClosedJaxpr.replace(jaxpr=..., consts=...) form. # TODO(dougalm): migrate callers and remove. jaxpr = kwargs.pop("jaxpr") consts = kwargs.pop("consts", None) if kwargs: raise ValueError(f"Unknown keyword arguments: {kwargs}") jaxpr = self if jaxpr is None else jaxpr consts = self.consts if consts is None else consts return Jaxpr(jaxpr, consts) debug_default = self.debug_info if (kwargs.get('invars', self.invars) != self.invars or kwargs.get('outvars', self.outvars) != self.outvars): debug_default = debug_default.with_unknown_names() # Replacing constvars invalidates the consts pairing, so unless new consts # are given explicitly the result has no consts attached. consts_default = () if "constvars" in kwargs else self.consts jaxpr = Jaxpr( constvars=kwargs.pop("constvars", self.constvars), invars=kwargs.pop("invars", self.invars), outvars=kwargs.pop("outvars", self.outvars), eqns=kwargs.pop("eqns", self.eqns), effects=kwargs.pop("effects", self.effects), debug_info=kwargs.pop("debug_info", debug_default), is_high=kwargs.pop("is_high", self.is_high), consts=kwargs.pop("consts", consts_default), ) if kwargs: raise ValueError(f"Unknown keyword arguments: {kwargs}") return jaxpr weakref_cache_key_types.add(Jaxpr) def _shift_arg_names(dbg: DebugInfo, delta: int) -> DebugInfo: # Keep debug arg_names aligned with `invars` as the const/invar boundary # moves: delta > 0 exposes that many constvars as invars (pad with unnamed # entries), delta < 0 turns that many leading invars into constvars. if not delta or dbg.arg_names is None: return dbg if delta > 0: return dbg._replace(arg_names=("",) * delta + tuple(dbg.arg_names)) return dbg._replace(arg_names=tuple(dbg.arg_names[-delta:])) def join_effects(*effects: Effects) -> Effects: return set().union(*effects) if effects else no_effects def jaxprs_in_params(params) -> Iterator[Jaxpr]: for val in params.values(): vals = val if isinstance(val, tuple) else (val,) for v in vals: if isinstance(v, Jaxpr): yield v def subjaxprs(jaxpr: Jaxpr) -> Iterator[Jaxpr]: """Generator for all subjaxprs found in the params of jaxpr.eqns. Does not descend recursively into the found subjaxprs. """ for eqn in jaxpr.eqns: yield from jaxprs_in_params(eqn.params) @curry def jaxpr_as_fun(closed_jaxpr: Jaxpr, *args): # TODO(dougalm): remove this hack when we add contexts to jaxpr. # debug_nans is sometimes disabled locally at the traceable level by ops that # work with nans internally, like jnp.var. The right thing to do is to add # contexts to our jaxpr representation so that we can capture these local # context modifications. In the meantime, disabling the checks when we # round-trip prevents those ops producing spurious errors. with config.debug_nans(False): return eval_jaxpr(closed_jaxpr, closed_jaxpr.consts, *args) # This context manager is fairly hot, because it is frequently called for every # jaxpr equation. # This context manager is implemented as a class with explicit __enter__ and # __exit__ methods since a @contextlib.contextmanager is significantly slower. # We also in effect fuse four other context managers into one, mostly to # save allocations. class JaxprEqnContextManager: __slots__ = ['context', 'prev_compute_type', 'prev_threefry_partitionable', 'prev_xla_metadata', 'prev_abstract_mesh', 'prev_layout_mode', 'prev_remove_size_one_mesh_axis'] def __init__(self, context): self.context = context def __enter__(self): self.prev_xla_metadata = config.xla_metadata_context_manager.get_local() if (self.context.xla_metadata and (self.prev_xla_metadata is None or self.prev_xla_metadata is config_ext.unset or self.prev_xla_metadata.val != self.context.xla_metadata)): updated = update_metadata( self.prev_xla_metadata, self.context.xla_metadata) config.xla_metadata_context_manager.set_local(updated) self.prev_threefry_partitionable = config.threefry_partitionable.swap_local( self.context.threefry_partitionable) self.prev_compute_type = config.compute_on_context_manager.swap_local( self.context.compute_type) self.prev_abstract_mesh = mesh_lib.abstract_mesh_context_manager.swap_local( self.context.cur_abstract_mesh) self.prev_layout_mode = layout_tracing_mode.swap_local( self.context.cur_layout_mode) self.prev_remove_size_one_mesh_axis = config.remove_size_one_mesh_axis_from_type.swap_local( self.context.remove_size_one_mesh_axis) def __exit__(self, exc_type, exc_value, traceback): config.xla_metadata_context_manager.set_local(self.prev_xla_metadata) config.threefry_partitionable.set_local(self.prev_threefry_partitionable) config.compute_on_context_manager.set_local(self.prev_compute_type) mesh_lib.abstract_mesh_context_manager.set_local(self.prev_abstract_mesh) layout_tracing_mode.set_local(self.prev_layout_mode) config.remove_size_one_mesh_axis_from_type.set_local(self.prev_remove_size_one_mesh_axis) @immutable class JaxprEqnContext: __slots__ = ['compute_type', 'threefry_partitionable', 'cur_abstract_mesh', 'cur_layout_mode', 'remove_size_one_mesh_axis', 'xla_metadata', 'configs', '__weakref__'] compute_type: str | None threefry_partitionable: bool xla_metadata: dict[str, Any] | None cur_abstract_mesh: mesh_lib.AbstractMesh cur_layout_mode: LayoutMode remove_size_one_mesh_axis: bool @staticmethod @weak_value_interner def _create(compute_type, threefry_partitionable, cur_abstract_mesh, cur_layout_mode, remove_size_one_mesh_axis, xla_metadata): obj = object.__new__(JaxprEqnContext) object.__setattr__(obj, 'compute_type', compute_type) object.__setattr__(obj, 'threefry_partitionable', threefry_partitionable) object.__setattr__(obj, 'cur_abstract_mesh', cur_abstract_mesh) object.__setattr__(obj, 'cur_layout_mode', cur_layout_mode) object.__setattr__(obj, 'remove_size_one_mesh_axis', remove_size_one_mesh_axis) object.__setattr__(obj, 'xla_metadata', None if xla_metadata is None else dict(xla_metadata)) return obj def __new__(cls): compute_type = config.compute_on_context_manager.value threefry_partitionable = config.threefry_partitionable.value cur_abstract_mesh = mesh_lib.get_abstract_mesh() cur_layout_mode = get_layout_mode() remove_size_one_mesh_axis = config.remove_size_one_mesh_axis_from_type.value xla_metadata = (None if (xm := current_xla_metadata()) is None else tuple(sorted(xm.items()))) return JaxprEqnContext._create( compute_type, threefry_partitionable, cur_abstract_mesh, cur_layout_mode, remove_size_one_mesh_axis, xla_metadata) # No __eq__ or __hash__: interned classes use object identity. @property def manager(self): return JaxprEqnContextManager(self) def __repr__(self): return (f"JaxprEqnContext(compute_type={self.compute_type}, " f"threefry_partitionable={self.threefry_partitionable}, " f"cur_abstract_mesh={self.cur_abstract_mesh}, " f"cur_layout_mode={self.cur_layout_mode}, " f"remove_size_one_mesh_axis={self.remove_size_one_mesh_axis}, " f"xla_metadata={self.xla_metadata})") @cache() # Everything in the context is a trace cache key also. def current_jaxpr_eqn_context(): return JaxprEqnContext() class JaxprEqn: invars: list[Atom] outvars: list[Var] primitive: Primitive params: dict[str, Any] effects: Effects # The source_info.name_stack is always relative to the enclosing jaxpr (only) # and does not include any name context from the caller of the jaxpr. A jaxpr # might have multiple callers, after all. # TODO(phawkins): update source_info.tracebacks to also be relative to the # enclosing jaxpr. source_info: source_info_util.SourceInfo ctx: JaxprEqnContext # It's slightly faster to use a class with __slots__ than a NamedTuple. __slots__ = ['invars', 'outvars', 'primitive', 'params', 'effects', 'source_info', 'ctx'] def __init__(self, invars, outvars, primitive, params, effs, source_info, ctx): self.invars = invars self.outvars = outvars self.primitive = primitive self.params = params self.effects = effs self.source_info = source_info self.ctx = ctx def __repr__(self): return str(pp_eqn(self, JaxprPpContext(), JaxprPpSettings())).rstrip() def replace( self, invars: list[Atom] | None = None, outvars: list[Var] | None = None, primitive: Primitive | None = None, params: dict[str, Any] | None = None, effects: Effects | None = None, source_info: source_info_util.SourceInfo | None = None, ctx: JaxprEqnContext | None = None ): return JaxprEqn( self.invars if invars is None else invars, self.outvars if outvars is None else outvars, self.primitive if primitive is None else primitive, self.params if params is None else params, self.effects if effects is None else effects, self.source_info if source_info is None else source_info, self.ctx if ctx is None else ctx, ) # TODO(mattjj): call typecheck rules here, so we don't form bad eqns def new_jaxpr_eqn(invars, outvars, primitive, params, effects, source_info=None, ctx=None) -> JaxprEqn: source_info = source_info or source_info_util.new_source_info() ctx = ctx or current_jaxpr_eqn_context() effects = resolve_input_effects(effects, invars) if config.enable_checks.value: assert all(isinstance(x, (Var, Literal)) for x in invars) assert all(isinstance(v, Var) for v in outvars) return JaxprEqn(invars, outvars, primitive, params, effects, source_info, ctx) def resolve_input_effects(effs, invars) -> Effects: if not any(isinstance(e, effects.JaxprInputEffect) and isinstance(e.input, int) for e in effs): return effs out_effs = set() for eff in effs: if isinstance(eff, effects.JaxprInputEffect) and isinstance(eff.input, int): invar = invars[eff.input] if isinstance(invar, Literal): continue eff = eff.replace(invar) out_effs.add(eff) return out_effs class Var: __slots__ = ["aval"] aval: AbstractValue def __init__(self, aval: AbstractValue): assert isinstance(aval, AbstractValue), aval self.aval = aval def __repr__(self): return f'Var(id={id(self)}):{self.aval.str_short()}' def pretty_print(self, context: JaxprPpContext, *, print_dtype: bool = True): del print_dtype # unused return f"{context.var_names[self]}" gensym = lambda: Var # In a jaxpr, `dropvar` can appear in place of a bound variable to indicate that # the assignment is dropped, i.e. that an expression's output value will never # be read. In that sense, `dropvar` is not a variable, but it is convenient to # treat it as a special case of one. Its `aval` is similarly inexact. class DropVar(Var): def __init__(self, aval: AbstractValue): super().__init__(aval) def __repr__(self): return '_' def pretty_print(self, context: JaxprPpContext, *, print_dtype: bool = True): del context, print_dtype # unused return '_' @final class Literal: # See https://docs.jax.dev/en/latest/internals/constants.html __slots__ = ["val", "aval"] val: Any aval: AbstractValue def __init__(self, val, aval): self.val = val self.aval = aval @property def hash(self): try: return hash(self.val) except TypeError: if type(self.val) in literalable_types: try: return hash((self.val.item(), self.val.dtype)) except (TypeError, AttributeError, ValueError): return None __hash__ = None def pretty_print(self, context: JaxprPpContext, *, print_dtype: bool = True): del context # unused dtype = getattr(self.aval, 'dtype', None) if not np.shape(self.val): val_str = str(np.asarray(self.val).item()) else: val_str = "[...]" if print_dtype and dtype: return f'{val_str}:{self.aval.str_short(short_dtypes=True)}' else: return val_str def __repr__(self): return f'Literal({self.val})' # The types of constants that can be used with core.Literal. Other constants # end up as `constvars`. literalable_types: set[type] = set() literalable_scalar_types: set[type] = set() def is_literalable(x: Any, for_ad: bool = False) -> bool: x_type = type(x) # Faster path for scalar types, which avoids an np.ndarray conversion. if x_type in literalable_scalar_types: return True # See https://docs.jax.dev/en/latest/internals/constants.html # for_ad: we want to preserve under AD if config.use_simplified_jaxpr_constants.value: from jax._src.array import ArrayImpl # pyrefly: ignore[missing-import] do_lit_array = not for_ad if isinstance(x, ArrayImpl): return do_lit_array else: do_lit_array = False for t in x_type.__mro__: if t in literalable_types: return (do_lit_array or not np.ndim(x)) return False def is_hoistable(v: Literal) -> bool: return (np.ndim(v.val) > 0 and getattr(v.val, "nbytes", 4) > config.embedded_constants_max_bytes.value) @partial(weakref_lru_cache, trace_context_in_key=False) def jaxpr_const_args(jaxpr: Jaxpr) -> list[tuple[ArrayLike, AbstractValue]]: # The non-scalar constants in core.Literal, in the entire Jaxpr, # uniquified by id. These will be hoisted as const arguments to the functions # in which they appear. # See https://docs.jax.dev/en/latest/internals/constants.html if not config.use_simplified_jaxpr_constants.value: return [] consts_by_id: dict[int, tuple[ArrayLike, AbstractValue]] = {} for v in jaxpr.outvars: if type(v) is Literal and is_hoistable(v): consts_by_id[id(v)] = (v.val, v.aval) for eqn in jaxpr.eqns: for v in eqn.invars: if type(v) is Literal and is_hoistable(v): consts_by_id[id(v)] = (v.val, v.aval) consts_by_id.update({id(v_aval[0]): v_aval for v_aval in eqn_params_const_args(eqn.params)}) return list(consts_by_id.values()) def eqn_params_const_args(params) -> list[tuple[ArrayLike, AbstractValue]]: consts_by_id: dict[int, tuple[ArrayLike, AbstractValue]] = {} for j in jaxprs_in_params(params): consts_by_id.update( {id(v_aval[0]): v_aval for v_aval in jaxpr_const_args(j)} ) return list(consts_by_id.values()) Atom = Var | Literal class Primitive: name: str # set for multi-output primitives. multiple_results: bool = False # set for call primitives processed in final style. call_primitive: bool = False # set for ref primitives ref_primitive: bool = False # set for primitives that can skip canonicalization of values skip_canonicalization: bool = False # set for primitives that allocate references ref_allocating: bool = False is_effectful = None def __init__(self, name: str): self.name = name def __repr__(self): return f'{self.name}' def bind(self, *args, **params): canonical_args = [] avals = [] for i, arg in enumerate(args): try: c_arg = dtypes.canonicalize_value(arg) aval = typeof(c_arg) except TypeError as e: raise TypeError( f"Error interpreting argument to {self} as a JAX value." f" The problematic value is of type {type(arg)} and was passed to" f" {self} at position {i}.\n" ) from e if (not self.skip_canonicalization and isinstance(aval, ShapedArray) and not aval.sharding.mesh.empty): cur_mesh = mesh_lib.get_abstract_mesh() if cur_mesh != aval.sharding.mesh: # TODO(yashkatariya): Casting to Explicit is not yet allowed. Maybe we # need cast_and_slice_p for it since shape might change? # Atleast 1 mesh axis should be Manual and all other axes should be # Manual or Auto to allow casting. if cur_mesh._any_axis_manual and cur_mesh._are_all_axes_auto_or_manual: if aval.sharding.mesh.are_all_axes_auto: from jax._src.pjit import reshard # pyrefly: ignore[missing-import] c_arg = reshard(c_arg, NamedSharding(cur_mesh, P(*[None] * aval.ndim))) aval = typeof(c_arg) elif aval.sharding.mesh._any_axis_explicit: raise NotImplementedError( "Closing over inputs to shard_map where the input is sharded " "on `Explicit` axes is not implemented. As a workaround, " "please pass those inputs as an argument to shard_map. Got " f"input with shape {aval.str_short(True, True)}") if isinstance(c_arg, Tracer) and not c_arg._trace.is_valid(): raise escaped_tracer_error(c_arg) canonical_args.append(c_arg) avals.append(aval) args = canonical_args # This is equivalent to "with take_current_trace()", but the bind() code # is called frequently and it's slightly faster to avoid using a context # manager object. prev_trace = trace_ctx.trace trace_ctx.set_trace(None) try: return self.bind_with_trace(prev_trace, args, avals, params) finally: trace_ctx.set_trace(prev_trace) def bind_with_trace(self, trace, args, avals, params, /): if self.is_high(*avals, **params) and trace.requires_low: with set_current_trace(trace): return self.to_lojax(*args, **params) return trace.process_primitive(self, args, params) def def_impl(self, impl): self.impl = impl return impl def def_abstract_eval(self, abstract_eval): self.abstract_eval = _effect_free_abstract_eval(abstract_eval) return abstract_eval def def_effectful_abstract_eval(self, effectful_abstract_eval): self.abstract_eval = effectful_abstract_eval return effectful_abstract_eval def def_effectful_abstract_eval2(self, abstract_eval): self.abstract_eval = _generic_effectful_abstract_eval(abstract_eval, self) return abstract_eval def def_bind_with_trace(self, bind_with_trace): self.bind_with_trace = bind_with_trace return bind_with_trace def impl(self, *args, **params): raise NotImplementedError("Evaluation rule for '{}' not implemented" .format(self.name)) def abstract_eval(self, *args, **params): raise NotImplementedError("Abstract evaluation for '{}' not implemented" .format(self.name)) def get_bind_params(self, params): return params def to_lojax(self, *args, **params): raise NotImplementedError( f"Primitive '{self.name}' has is_high=True but no to_lojax " f"lowering rule.") def is_high(self, *avals, **params) -> bool: for v in params.values(): if isinstance(v, Jaxpr) and v.is_high: return True return False def _effect_free_abstract_eval(abstract_eval): def abstract_eval_(*args, **kwargs): return abstract_eval(*args, **kwargs), no_effects return abstract_eval_ @dataclass(frozen=True, slots=True) class GenericEffect(Effect): prim: Primitive effects.lowerable_effects.add_type(GenericEffect) effects.control_flow_allowed_effects.add_type(GenericEffect) effects.custom_derivatives_allowed_effects.add_type(GenericEffect) def _generic_effectful_abstract_eval(abstract_eval, prim): def abstract_eval_(*args, **kwargs): return abstract_eval(*args, **kwargs), {GenericEffect(prim)} return abstract_eval_ # -------------------- lifting -------------------- def eval_jaxpr(jaxpr: Jaxpr, consts, *args, propagate_source_info=True) -> list[Any]: def read(v: Atom) -> Any: return v.val if isinstance(v, Literal) else env[v] def write(v: Var, val: Any) -> None: if config.enable_checks.value: assert typecheck(v.aval, val), (v.aval, typeof(val), val) env[v] = val env: dict[Var, Any] = {} foreach(write, jaxpr.all_invars, [*consts, *args]) lu = last_used(jaxpr) for eqn in jaxpr.eqns: bind_params = eqn.primitive.get_bind_params(eqn.params) name_stack = source_info_util.current_name_stack() + eqn.source_info.name_stack traceback = eqn.source_info.traceback if propagate_source_info else None with (source_info_util.user_context(traceback, name_stack=name_stack), eqn.ctx.manager): ans = eqn.primitive.bind(*map(read, eqn.invars), **bind_params) if eqn.primitive.multiple_results: foreach(write, eqn.outvars, ans) else: write(eqn.outvars[0], ans) clean_up_dead_vars(eqn, env, lu) return map(read, jaxpr.outvars) def check_avals_context_mesh(avals, prim_name): cur_mesh = mesh_lib.get_abstract_mesh() for a in avals: if not isinstance(a.memory_space, MemorySpace): raise TypeError( f"Primitive {prim_name} got aval {a} with unknown memory_space type:" f" {type(a.memory_space)}") if cur_mesh.empty or a.sharding.mesh.empty: continue # avals can have meshes with different axis_names so allow that in # full auto mode. if a.sharding.mesh.are_all_axes_auto and cur_mesh.are_all_axes_auto: continue if a.sharding.mesh != cur_mesh: raise ValueError( f"For primitive {prim_name}, context mesh {cur_mesh} should match" f" the aval mesh {a.sharding.mesh} for shape {a.str_short()}. This" " error occurs at source: " f" {source_info_util.summarize(source_info_util.current())}") # -------------------- tracing -------------------- class Trace: __slots__ = ("__weakref__", "_invalidated", "_weakref", "requires_low") def __init__(self): self._invalidated = False # We frequently need a weakref to a trace, so let's precompute one. self._weakref = weakref.ref(self) self.requires_low = True def stage_value(self, val): """Lifts a value into a trace. Semantically equivalent to calling process_primitive on an identity primitive, but may avoid, e.g., constructing a jaxpr equation.""" raise NotImplementedError("must override") def process_primitive(self, primitive, tracers, params, /): raise NotImplementedError("must override") def invalidate(self): self._invalidated = True def is_valid(self): return not self._invalidated def __repr__(self): return f'{self.__class__.__name__}' def process_custom_jvp_call(self, primitive, fun, jvp, tracers, /, *, symbolic_zeros): msg = (f"{type(self)} must override process_custom_jvp_call " "to handle custom_jvp primitives") raise NotImplementedError(msg) def process_custom_vjp_call(self, primitive, fun, fwd, bwd, tracers, /, *, out_trees, symbolic_zeros): msg = (f"{type(self)} must override process_custom_vjp_call " "to handle custom_vjp primitives") raise NotImplementedError(msg) # TODO(dougalm): deprecate/delete def full_raise(self, x): return x # TODO(dougalm): deprecate/delete @property def main(self): return getattr(self, "tag", None) def escaped_tracer_error(tracer, detail=None): num_frames = _TRACER_ERROR_NUM_TRACEBACK_FRAMES.value msg = ('Encountered an unexpected tracer. A function transformed by JAX ' 'had a side effect, allowing for a reference to an intermediate value ' f'with type {tracer.aval.str_short()} wrapped in a ' f'{type(tracer).__name__} to escape the scope of the transformation.\n' 'JAX transformations require that functions explicitly return their ' 'outputs, and disallow saving intermediate values to global state.') dbg = getattr(tracer, '_debug_info', None) if dbg is not None: msg += ('\nThe function being traced when the value leaked was ' f'{dbg.func_src_info} traced for {dbg.traced_for}.') line_info = getattr(tracer, '_line_info', None) if line_info is not None: divider = '\n' + '-'*30 + '\n' msg += divider msg += ('The leaked intermediate value was created on line ' f'{source_info_util.summarize(line_info)}. ') msg += divider if num_frames > 0: msg += (f'When the value was created, the final {num_frames} stack ' 'frames (most recent last) excluding JAX-internal frames were:') msg += divider + source_info_util.summarize( line_info, num_frames=num_frames) + divider msg += ('\nTo catch the leak earlier, try setting the environment variable ' 'JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context ' 'manager.') if detail: msg += f'Detail: {detail}' return UnexpectedTracerError(msg) def check_scalar_conversion(arr: Array): if arr.ndim > 0: raise TypeError("Only scalar arrays can be converted to Python scalars; " f"got {arr.ndim=}") def check_integer_conversion(arr: Array): if not (arr.shape == () and dtypes.issubdtype(arr.dtype, np.integer)): raise TypeError("Only integer scalar arrays can be converted to a scalar index.") def check_bool_conversion(arr: Array): if arr.size == 0: raise ValueError("The truth value of an empty array is ambiguous. Use" " `array.size > 0` to check that an array is not empty.") if arr.size > 1: raise ValueError("The truth value of an array with more than one element" " is ambiguous. Use a.any() or a.all()") pytype_aval_mappings: dict[type, Callable[[Any], AbstractValue]] = {} def _str_abstractify(x): raise TypeError(f"Argument '{x}' of type {type(x)} is not a valid JAX type") pytype_aval_mappings[str] = _str_abstractify def _aval_property(name): return property(lambda self: getattr(self.aval, name)) if TYPE_CHECKING: # We want Python type checkers to accept `some_tracer: jax.Array`, even though # tracers can represent non-arrays. That is, ideally we would only accept that # annotation when the Tracer instance has a ShapedArray aval, but we can't # decide that at Python type checking time. So instead we're overly permissive # and allow all Tracer instances to typecheck against a jax.Array annotation. TracerBase = Array TracerMeta = StrictABCMeta else: class TracerBase: __slots__ = () TracerMeta = type class Tracer[TraceType: Trace](TracerBase, metaclass=TracerMeta): __array_priority__ = 1000 __slots__ = ['__weakref__', '_trace', '_line_info', 'aval'] __hash__ = None _trace: TraceType _line_info: source_info_util.SourceInfo | None aval: AbstractValue dtype = _aval_property('dtype') ndim = _aval_property('ndim') size = _aval_property('size') shape = _aval_property('shape') # dimension_as_value is frequently accessed, and we explicitly define it as # None to avoid hitting the __getattr__ path, which constructs an error # message (and is therefore slow). dimension_as_value = None # We define __jax_array__ as a property to delegate to self.aval.__jax_array__ # if it exists (e.g., for avals like Flax NNX variables). # This avoids hitting the slow __getattr__ path for tracers that don't have it. @property def __jax_array__(self): m = getattr(self.aval, '__jax_array__', None) if m is not None: return lambda: m.fun(self) return None def __init__(self, trace: TraceType, aval: AbstractValue): self._trace = trace self.aval = aval def _error_repr(self): if self.aval is None: return f"traced array with aval {self.aval}" return f"traced array with shape {self.aval.str_short()}" def __array__(self, *args, **kw): raise TracerArrayConversionError(self) # helper for isinstance(tracer, jax.Array), here to avoid circular imports def _is_traced_array(self): return isinstance(self.aval, ShapedArray) def __dlpack__(self, *args, **kw): raise ConcretizationTypeError(self, f"The __dlpack__() method was called on {self._error_repr()}." f"{self._origin_msg()}") def tolist(self): raise ConcretizationTypeError(self, f"The tolist() method was called on {self._error_repr()}." f"{self._origin_msg()}") def tobytes(self, order="C"): del order raise ConcretizationTypeError(self, f"The tobytes() method was called on {self._error_repr()}." f"{self._origin_msg()}") # TODO(dougalm): deprecate/delete def full_lower(self): return self def __iter__(self): if not hasattr(self.aval, "_iter"): raise TypeError(f"Value of type {type(self)} is not iterable.") return iter(self.aval._iter(self)) def __reversed__(self): return iter(self[::-1]) def __len__(self): if not hasattr(self.aval, "_len"): raise TypeError(f"Value of type {type(self)} has no length.") return self.aval._len(self) def to_concrete_value(self): # Should return the concrete value if there is one, or else None. return None @property def sharding(self): # This attribute is part of the jax.Array API, but only defined on concrete arrays. # Raising a ConcretizationTypeError would make sense, but for backward compatibility # we raise an AttributeError so that hasattr() and getattr() work as expected. raise AttributeError( f"The 'sharding' attribute is not available on {self._error_repr()}.") @property def memory_space(self): # we raise ValueError instead of AttributeError to avoid silent bugs due # to hasattr/getattr returning False instead of showing the error with the # correct way to access the memory space. raise ValueError( "The 'memory_space' attribute is not available on " f"{self._error_repr()}. Use `jax.typeof(x).memory_space` instead.") @property def committed(self): raise ConcretizationTypeError( self, f"The 'committed' attribute is not available on {self._error_repr()}." f"{self._origin_msg()}") @property def device(self): # This attribute is part of the jax.Array API, but only defined on concrete arrays. # Raising a ConcretizationTypeError would make sense, but for backward compatibility # we raise an AttributeError so that hasattr() and getattr() work as expected. raise AttributeError( f"The 'device' attribute is not available on {self._error_repr()}.") @property def addressable_shards(self): raise ConcretizationTypeError(self, f"The 'addressable_shards' attribute is not available on {self._error_repr()}." f"{self._origin_msg()}") @property def at(self): if not hasattr(self.aval, "at"): raise TypeError(f"Value of type {type(self)} does not support at().") return self.aval.at.fget(self) def get_referent(self) -> Any: return self # Override for object equivalence checking def __bool__(self): if is_concrete(self): return bool(self.to_concrete_value()) check_bool_conversion(self) if not hasattr(self.aval, "_bool"): raise TypeError(f"Value of type {type(self)} is not convertible to boolean.") return self.aval._bool(self) def __int__(self): if is_concrete(self): return int(self.to_concrete_value()) check_scalar_conversion(self) if not hasattr(self.aval, "_int"): raise TypeError(f"Value of type {type(self)} is not convertible to integer.") return self.aval._int(self) def __float__(self): check_scalar_conversion(self) if not hasattr(self.aval, "_float"): raise TypeError(f"Value of type {type(self)} is not convertible to float.") return self.aval._float(self) def __complex__(self): check_scalar_conversion(self) if not hasattr(self.aval, "_complex"): raise TypeError(f"Value of type {type(self)} is not convertible to complex.") return self.aval._complex(self) def __hex__(self): if is_concrete(self): return hex(self.to_concrete_value()) check_integer_conversion(self) if not hasattr(self.aval, "_hex"): raise TypeError(f"Value of type {type(self)} is not convertible to hex.") return self.aval._hex(self) def __oct__(self): if is_concrete(self): return oct(self.to_concrete_value()) check_integer_conversion(self) if not hasattr(self.aval, "_oct"): raise TypeError(f"Value of type {type(self)} is not convertible to oct.") return self.aval._oct(self) def __index__(self): if is_concrete(self): return operator.index(self.to_concrete_value()) check_integer_conversion(self) if not hasattr(self.aval, "_index"): raise TypeError(f"Value of type {type(self)} is not convertible to integer index.") return self.aval._index(self) # raises a useful error on attempts to pickle a Tracer. def __reduce__(self): raise ConcretizationTypeError( self, ("The error occurred in the __reduce__ method, which may " "indicate an attempt to serialize/pickle a traced value.")) # raises the better error message from ShapedArray def __setitem__(self, key, value): if not hasattr(self.aval, "_setitem"): raise TypeError(f"Value of type {type(self)} is not indexable.") return self.aval._setitem(self, key, value) # NumPy also only looks up special methods on classes. def __array_module__(self, types): if not hasattr(self.aval, "_array_module"): raise TypeError(f"Value of type {type(self)} is not compatible with the Array API.") return self.aval._array_module(self, types) def __getattr__(self, name): # if the aval property raises an AttributeError, gets caught here assert not config.enable_checks.value or name != "aval" # These must raise AttributeError in the base class for backward compatibility. # TODO(jakevdp): can we change this and make them raise NotImplementedError instead? if name in ["block_until_ready", "copy_to_host_async"]: raise AttributeError( f"The '{name}' method is not available on {self._error_repr()}." f"{self._origin_msg()}") if name == 'sharding': raise AttributeError( f"The 'sharding' attribute is not available on {self._error_repr()}. " "To query sharding information on tracers, use `jax.typeof(x)`.") try: attr = getattr(self.aval, name) except AttributeError as err: raise AttributeError( f"{self.__class__.__name__} has no attribute {name}" ) from err else: t = type(attr) if t is aval_property: return attr.fget(self) elif t is aval_method: return types.MethodType(attr.fun, self) else: return attr def _short_repr(self) -> str: return f'{self.__class__.__name__}<{self.aval}>' def _pretty_print(self, verbose: bool = False) -> pp.Doc: if not verbose: return pp.text(self._short_repr()) base = pp.text(f'Traced<{self.aval}>with<{self._trace}>') contents = [(name, attr._pretty_print() if isinstance(attr, Tracer) else pp.text(repr(attr))) for name, attr in self._contents()] if contents: base = pp.group(pp.nest(2, pp.concat([ base, pp.text(' with'), pp.brk(), pp.join(pp.brk(), [ pp.text(f'{name} = ') + pp_payload for name, pp_payload in contents]) ]))) return base def __repr__(self): return self._pretty_print(verbose=False).format() def _contents(self): try: return [(name, getattr(self, name)) for name in self.__slots__] except AttributeError: return () def _origin_msg(self) -> str: return "" # Methods that are only valid for materialized arrays def addressable_data(self, index): raise ConcretizationTypeError(self, f"The addressable_data() method was called on {self._error_repr()}." f"{self._origin_msg()}") def delete(self): raise ConcretizationTypeError(self, f"The delete() method was called on {self._error_repr()}." f"{self._origin_msg()}") def devices(self): raise ConcretizationTypeError(self, f"The devices() method was called on {self._error_repr()}." f"{self._origin_msg()}") @property def global_shards(self): raise ConcretizationTypeError(self, f"The global_shards property was called on {self._error_repr()}." f"{self._origin_msg()}") def is_deleted(self): raise ConcretizationTypeError(self, f"The is_deleted() method was called on {self._error_repr()}." f"{self._origin_msg()}") @property def is_fully_addressable(self): raise ConcretizationTypeError(self, f"The is_fully_addressable property was called on {self._error_repr()}." f"{self._origin_msg()}") @property def is_fully_replicated(self): raise ConcretizationTypeError(self, f"The is_fully_replicated property was called on {self._error_repr()}." f"{self._origin_msg()}") def on_device_size_in_bytes(self): raise ConcretizationTypeError(self, f"The on_device_size_in_bytes() method was called on {self._error_repr()}." f"{self._origin_msg()}") @property def traceback(self): raise ConcretizationTypeError(self, f"The traceback property was called on {self._error_repr()}." f"{self._origin_msg()}") def unsafe_buffer_pointer(self): raise ConcretizationTypeError(self, f"The unsafe_buffer_pointer() method was called on {self._error_repr()}." f"{self._origin_msg()}") _jax.set_tracer_class(Tracer) # these can be used to set up forwarding of properties and instance methods from # Tracer instances to the underlying avals aval_property = namedtuple("aval_property", ["fget"]) aval_method = namedtuple("aval_method", ["fun"]) pytype_aval_mappings[Tracer] = lambda x: x.aval dtypes.register_canonicalize_value_handler(Tracer, None) def check_eval_args(args): for arg in args: if isinstance(arg, Tracer): raise escaped_tracer_error(arg) stage_p = Primitive('stage') class EvalTrace(Trace): def stage_value(self, val): if isinstance(val, Array): return val return self.process_primitive(stage_p, [val], {}) def process_primitive(self, primitive, args, params, /): with set_current_trace(self): if config.debug_key_reuse.value: from jax.experimental.key_reuse._core import call_impl_with_key_reuse_checks # pyrefly: ignore[missing-import] return call_impl_with_key_reuse_checks(primitive, primitive.impl, *args, **params) else: # TODO(dougalm): delete. this shouldn't be necessary args = map(full_lower, args) check_eval_args(args) return primitive.impl(*args, **params) def process_custom_jvp_call(self, primitive, fun, jvp, tracers, /, **_): del primitive, jvp, _ # Unused. with set_current_trace(self): return fun.call_wrapped(*tracers) def process_custom_vjp_call(self, primitive, fun, fwd, bwd, tracers, /, **_): del primitive, fwd, bwd, _ # Unused. with set_current_trace(self): return fun.call_wrapped(*tracers) class TraceTag: # TODO: this works for surprisingly subtle reasons. Function transformations # like `jvp_subtrace` are parameterized by a tag that identifies the set of # pre-existing tracers we want to unpack during the transformation. A function # defined in an outer scope can't have any closed-over traces, so the tag is # irrelevant. A function defined in the current scope may have closed-over # traces, but the tag will never change so we'll never get a spurious cache # hit. The plan is to do away with `lu.cache` altogether, and use a simpler # caching scheme that only caches top-level functions. Then we can remove this # hack. def __hash__(self): return hash(TraceTag) def __eq__(self, other): return isinstance(other, TraceTag) ParamDict = dict[str, Any] AxisName = Hashable no_axis_name = object() @immutable class AxisEnv: __slots__ = ('axis_sizes', 'spmd_axis_names', 'explicit_mesh_axis_names', '__weakref__') axis_sizes : FrozenDict[AxisName, int] spmd_axis_names : frozenset[AxisName] explicit_mesh_axis_names: frozenset[AxisName] @staticmethod @weak_value_interner def _create(axis_sizes, spmd_axis_names, explicit_mesh_axis_names): obj = object.__new__(AxisEnv) object.__setattr__(obj, 'axis_sizes', axis_sizes) object.__setattr__(obj, 'spmd_axis_names', spmd_axis_names) object.__setattr__(obj, 'explicit_mesh_axis_names', explicit_mesh_axis_names) return obj def __new__(cls, axis_sizes=FrozenDict({}), spmd_axis_names=frozenset(), explicit_mesh_axis_names=frozenset()): return cls._create(axis_sizes, spmd_axis_names, explicit_mesh_axis_names) def axis_size(self, axis_name): if axis_name not in self.axis_sizes: raise NameError(f"unbound axis name: {axis_name}") else: return self.axis_sizes[axis_name] def axis_exists(self, axis_name): return axis_name in self.axis_sizes def axis_names(self): return tuple(k for k in self.axis_sizes) def pop_pure(self, axis_name): new_sizes = dict(self.axis_sizes) new_sizes.pop(axis_name) return AxisEnv(FrozenDict(new_sizes), self.spmd_axis_names, self.explicit_mesh_axis_names) def extend_pure(self, name_size_pairs): new_sizes = dict(self.axis_sizes) new_sizes.update((name, size) for name, size in name_size_pairs if name is not no_axis_name) return AxisEnv(FrozenDict(new_sizes), self.spmd_axis_names, self.explicit_mesh_axis_names) def add_spmd_axis_names(self, axis_names): new_spmd_axis_names = self.spmd_axis_names | frozenset(axis_names) return AxisEnv(self.axis_sizes, new_spmd_axis_names, self.explicit_mesh_axis_names) def add_explicit_mesh_axis_names(self, axis_names): new_ema = self.explicit_mesh_axis_names | frozenset(axis_names) return AxisEnv(self.axis_sizes, self.spmd_axis_names, new_ema) def remove_explicit_mesh_axis_names(self, axis_names): new_ema = self.explicit_mesh_axis_names - frozenset(axis_names) return AxisEnv(self.axis_sizes, self.spmd_axis_names, new_ema) eval_trace = EvalTrace() top_axis_env = AxisEnv(FrozenDict({}), frozenset(), frozenset()) # Weak reference to the trace state. This is included in, e.g., the jit key. trace_state = config_ext.Config[Any]( 'trace_state', eval_trace._weakref, include_in_jit_key=True) # A strong reference to the trace state. This should not be included in any # jit or cache keys, but we need a thread-local strong reference to ensure it # remains alive. trace_state_strong_ref = config_ext.Config[Any]( 'trace_state_strong_ref', eval_trace, include_in_jit_key=False, include_in_trace_context=False) axis_env_state = config_ext.Config[Any]( 'axis_env_state', top_axis_env, include_in_jit_key=True, include_in_trace_context=True, ) class TracingContext: __slots__ = () @staticmethod def reset(): trace_state.set_local(config_ext.unset) trace_state_strong_ref.set_local(config_ext.unset) axis_env_state.set_local(config_ext.unset) @property def trace(self): return trace_state_strong_ref.value @property def axis_env(self): return axis_env_state.value def is_top_level(self) -> bool: return self.trace is eval_trace and self.axis_env is top_axis_env def is_empty(self) -> bool: return self.trace is None def set_trace(self, trace): trace_state_strong_ref.set_local(trace) ts = trace._weakref if trace is not None else None trace_state.set_local(ts) def set_axis_env(self, axis_env): axis_env_state.set_local(axis_env) trace_ctx = TracingContext() class TakeCurrentTraceContextManager: __slots__ = ['prev'] def __enter__(self): self.prev = trace_ctx.trace trace_ctx.set_trace(None) return self.prev def __exit__(self, exc_type, exc_value, traceback): trace_ctx.set_trace(self.prev) take_current_trace = TakeCurrentTraceContextManager class SetCurrentTraceContextManager: __slots__ = ['trace', 'check_leaks', 'prev'] def __init__(self, trace, check_leaks=False): self.trace = trace self.check_leaks = check_leaks def __enter__(self): self.prev = trace_ctx.trace trace_ctx.set_trace(self.trace) def __exit__(self, exc_type, exc_value, traceback): trace_ctx.set_trace(self.prev) if self.check_leaks and config.check_tracer_leaks.value: self.trace.invalidate() trace_ref = self.trace._weakref del self.trace live_trace = trace_ref() if live_trace is not None: leaked_tracers = maybe_find_leaked_tracers(live_trace) if leaked_tracers: raise leaked_tracer_error("trace", live_trace, leaked_tracers) set_current_trace = SetCurrentTraceContextManager class ExtendAxisEnvNdContextManager: __slots__ = ['prev', 'name_size_pairs'] def __init__(self, name_size_pairs: Iterable[tuple[AxisName, int]]): self.name_size_pairs = name_size_pairs def __enter__(self): self.prev = trace_ctx.axis_env trace_ctx.set_axis_env(self.prev.extend_pure(self.name_size_pairs)) def __exit__(self, exc_type, exc_value, traceback): trace_ctx.set_axis_env(self.prev) extend_axis_env_nd = ExtendAxisEnvNdContextManager class AddSpmdAxisNamesContextManager: __slots__ = ['prev', 'axis_names'] def __init__(self, axis_names: AxisName | None): self.axis_names = axis_names def __enter__(self): self.prev = trace_ctx.axis_env if self.axis_names is not None: trace_ctx.set_axis_env(self.prev.add_spmd_axis_names(self.axis_names)) def __exit__(self, exc_type, exc_value, traceback): trace_ctx.set_axis_env(self.prev) add_spmd_axis_names = AddSpmdAxisNamesContextManager # TODO(yashkatariya): Remove this once vmap handles mesh contexts correctly. class AddExplicitMeshAxisNamesContextManager: __slots__ = ['prev', 'axis_names'] def __init__(self, axis_names: AxisName | None): self.axis_names = axis_names def __enter__(self): self.prev = trace_ctx.axis_env if self.axis_names is not None: trace_ctx.set_axis_env(self.prev.add_explicit_mesh_axis_names( self.axis_names)) def __exit__(self, exc_type, exc_value, traceback): trace_ctx.set_axis_env(self.prev) add_explicit_mesh_axis_names = AddExplicitMeshAxisNamesContextManager # TODO(yashkatariya): Remove this once vmap handles mesh contexts correctly. class RemoveExplicitMeshAxisNamesContextManager: __slots__ = ['prev', 'axis_names'] def __init__(self, axis_names: AxisName | None): self.axis_names = axis_names def __enter__(self): self.prev = trace_ctx.axis_env if self.axis_names is not None: trace_ctx.set_axis_env(self.prev.remove_explicit_mesh_axis_names( self.axis_names)) def __exit__(self, exc_type, exc_value, traceback): trace_ctx.set_axis_env(self.prev) remove_explicit_mesh_axis_names = RemoveExplicitMeshAxisNamesContextManager def get_axis_env(): return trace_ctx.axis_env def trace_state_clean() -> bool: return trace_ctx.is_empty() or trace_ctx.is_top_level() def reset_trace_state() -> bool: """Resets the global trace state and returns True if it was already clean.""" if not trace_ctx.is_top_level(): trace_ctx.reset() return False else: return True TRACER_LEAK_DEBUGGER_WARNING = """\ JAX check_tracer_leaks behavior can trigger false positives when used with a debugger. To avoid false positives and silence this warning, you can disable thread tracing using the following: import threading threading.current_thread().pydev_do_not_trace = True """ @contextmanager def ensure_no_leaks(trace:Trace): yield trace.invalidate() if config.check_tracer_leaks.value: trace_ref = trace._weakref del trace live_trace = trace_ref() if live_trace is not None: leaked_tracers = maybe_find_leaked_tracers(live_trace) if leaked_tracers: raise leaked_tracer_error("trace", live_trace, leaked_tracers) def maybe_find_leaked_tracers(trace: Trace) -> list[Tracer]: """Find the leaked tracers holding a reference to the Trace """ if not getattr(threading.current_thread(), 'pydev_do_not_trace', True): warnings.warn(TRACER_LEAK_DEBUGGER_WARNING) # Trigger garbage collection to filter out unreachable objects that are alive # only due to cyclical dependencies. (We don't care about unreachable leaked # tracers since they can't interact with user code and cause a problem.) gc.collect() tracers = list(filter(lambda x: isinstance(x, Tracer), gc.get_referrers(trace))) return tracers def leaked_tracer_error(name: str, t, tracers: list[Tracer]) -> Exception: assert tracers why = partial(_why_alive, {id(tracers)}) msgs = [] for tracer in tracers: # not a genexpr: it'd be gc-visible and self-report chain = why(tracer) label = f'<{type(tracer).__name__} {id(tracer)}>' chain += ''.join(f'\n{label} is referred to by {h}' for h in _held_in_frame_locals(tracer, {id(tracers)})) if not chain: chain = (f'\n{label} has no referrers visible to the gc module; it may ' 'be held by an object that does not cooperate with the garbage ' 'collector, such as one implemented in a C extension') msgs.append(f'{tracer}{tracer._origin_msg()}{chain}') return Exception(f'Leaked {name} {t}. Leaked tracer(s):\n\n' + '\n\n'.join(msgs) + '\n') def _held_in_frame_locals(x, ignore_ids: set[int]) -> list[str]: """Find live stack frames whose locals refer to (or contain) x. On CPython 3.11+, executing functions' frames are usually not gc-tracked objects, so references held by their locals are invisible to gc.get_referrers and hence to _why_alive. Walk the current stack directly instead. """ skip_codes = (leaked_tracer_error.__code__, _held_in_frame_locals.__code__) holders = [] frame = sys._getframe(1) while frame is not None: if frame.f_code not in skip_codes: for name, val in frame.f_locals.items(): if id(val) in ignore_ids: continue if val is x: via = '' elif _contains_ref(val, x): via = f', a {type(val).__name__} containing it,' else: continue code = frame.f_code holders.append(f"the local variable {name!r}{via} of the frame " f"{code.co_qualname} ({code.co_filename}:{frame.f_lineno})") frame = frame.f_back return holders def _contains_ref(val, x, depth: int = 0) -> bool: if depth >= 3: return False if isinstance(val, (list, tuple, set, frozenset)): return any(v is x or _contains_ref(v, x, depth + 1) for v in val) if isinstance(val, dict): return any(v is x or _contains_ref(v, x, depth + 1) for v in val.values()) return False def _why_alive(ignore_ids: set[int], x: Any) -> str: parents = lambda x: [r for r in gc.get_referrers(x) if id(r) not in ignore_ids] child, lines, seen = x, [], set() while (id(child) not in seen and type(child) is not types.ModuleType and parents(child)): parent = parents(child)[0] # just pick one parent # For namespaces (like modules and class instances) and closures, the # references may form a simple chain: e.g. instance refers to its own # __dict__ which refers to child, or function refers to its __closure__ # which refers to cells which refer to child. In these cases, we can provide # a more intuitive description by collapsing the chain into a single # parent->child jump. We do that by setting `parent` here to be a # grandparent (or great-grandparent) of `child`, and then handling that case # in _why_alive_container_info. See example: # https://github.com/jax-ml/jax/pull/13022#discussion_r1008456599 # To prevent this collapsing behavior, just comment out this code block. try: if (isinstance(parent, dict) and getattr(parents(parent)[0], '__dict__', None) is parents(child)[0]): parent = parents(parent)[0] elif type(parent) is types.CellType: parent = parents(parents(parent)[0])[0] except IndexError: pass # a referrer list can be empty, e.g. a container held only by a # live frame's local, since gc.get_referrers can't see live frames line = f'<{type(child).__name__} {id(child)}> is referred to by ' lines.append(line + _why_alive_container_info(parent, id(child))) seen.add(id(child)) child = parent return '\n' + '\n'.join(lines) if lines else '' def _why_alive_container_info(container, obj_id) -> str: name = f'<{type(container).__name__} {id(container)}>' if type(container) is types.ModuleType: name = getattr(container, '__name__', name) if type(container) is types.FunctionType: name_ = getattr(container, '__name__', '') closure = inspect.getclosurevars(container) keys = [k for k, v in dict(closure.nonlocals, **closure.globals).items() if id(v) == obj_id] if len(keys) == 1: return f'{name} ({name_}) closed-over variable {keys[0]}' elif len(keys) > 1: return (f'{name} in closed-over variables ' + ', '.join(map(repr, keys))) if hasattr(container, '__dict__'): keys = [k for k in vars(container) if id(vars(container)[k]) == obj_id] if len(keys) == 1: return f'{name}.{keys[0]}' elif len(keys) > 1: return f'{name} in vars ' + ', '.join(map(repr, keys)) if isinstance(container, (list, tuple)): idxs = [i for i, x in enumerate(container) if id(x) == obj_id] if len(idxs) == 1: return f'{name}[{idxs[0]}]' else: return f'{name} at indices ' + ', '.join(map(str, idxs)) if isinstance(container, dict): keys = [k for k in container if id(container[k]) == obj_id] if len(keys) == 1: return f'{name}[{keys[0]!r}]' else: return f'{name} at keys ' + ', '.join(map(repr, keys)) if isinstance(container, types.ModuleType): return f' named {container.__name__}' return name @contextmanager def ensure_compile_time_eval(): """Context manager to ensure evaluation at trace/compile time (or error). Some JAX APIs like :func:`jax.jit` and :func:`jax.lax.scan` involve staging, i.e., delaying the evaluation of numerical expressions (like :mod:`jax.numpy` function applications) so that instead of performing those computations eagerly while evaluating the corresponding Python expressions, their computation is carried out separately, e.g. after optimized compilation. But this delay can be undesirable. For example, numerical values might be needed to evaluate Python control flow and so their evaluation cannot be delayed. As another example, it may be beneficial to ensure compile time evaluation (or "constant folding") for performance reasons. This context manager ensures that JAX computations are evaluated eagerly. If eager evaluation is not possible, a ``ConcretizationTypeError`` is raised. Here's a contrived example:: import jax import jax.numpy as jnp @jax.jit def f(x): with jax.ensure_compile_time_eval(): y = jnp.sin(3.0) z = jnp.sin(y) z_positive = z > 0 if z_positive: # z_positive is usable in Python control flow return jnp.sin(x) else: return jnp.cos(x) Here's a real-world example from https://github.com/jax-ml/jax/issues/3974:: import jax import jax.numpy as jnp from jax import random @jax.jit def jax_fn(x): with jax.ensure_compile_time_eval(): y = random.randint(random.key(0), (1000,1000), 0, 100) y2 = y @ y x2 = jnp.sum(y2) * x return x2 A similar behavior can often be achieved simply by 'hoisting' the constant expression out of the corresponding staging API:: y = random.randint(random.key(0), (1000,1000), 0, 100) @jax.jit def jax_fn(x): y2 = y @ y x2 = jnp.sum(y2)*x return x2 But in some cases it can be more convenient to use this context manager. """ with config.eager_constant_folding(True): yield @contextmanager def eval_context(): with set_current_trace(eval_trace): yield # TODO(dougalm): deprecate/delete def full_lower(val): if isinstance(val, Tracer): return val.full_lower() else: return val def get_referent(x: Any) -> Any: return x.get_referent() if isinstance(x, Tracer) else x def same_referent(x: Any, y: Any) -> bool: return get_referent(x) is get_referent(y) def dedup_referents(itr: Iterable[Any]) -> list[Any]: return list({HashableWrapper(get_referent(x)):x for x in itr}.values()) def definitely_equal(x, y): if isinstance(x, Tracer) or isinstance(y, Tracer): return same_referent(x, y) elif x is y: return True try: return x == y except InconclusiveDimensionOperation: return False # -------------------- abstract values -------------------- class AbstractValue: __slots__: list[str] = [] @property def is_high(self) -> bool: return False def to_tangent_aval(self) -> AbstractValue: raise NotImplementedError("must override") def to_ct_aval(self) -> AbstractValue: raise NotImplementedError("must override") # TODO(dougalm): deprecate this alias def at_least_vspace(self): return self.to_tangent_aval() def __repr__(self): try: kv_pairs = (f'{k}={v}' for k, v in self.__dict__.items()) return '{}({})'.format(self.__class__.__name__, ','.join(kv_pairs)) except AttributeError: return self.__class__.__name__ def update_weak_type(self, weak_type): return self def update_manual_axis_type(self, mat): return self def strip_weak_type(self) -> AbstractValue: return self.update_weak_type(False) def normalize(self) -> AbstractValue: return self.strip_weak_type() def update(self, **kwargs): raise NotImplementedError("must override") def lo_ty(self): return [self] def lower_val(self, val, /): return [val] def raise_val(self, *vals): val, = vals return val def str_short(self, short_dtypes=False, mesh_axis_types=False): return str(self) def dec_rank(self, size, spec): return mapped_aval(size, spec, self) def inc_rank(self, size, spec): return unmapped_aval(size, spec, self) def leading_axis_spec(self): return 0 def shard(self, mesh, manual_axes, check_vma, spec): return shard_aval(mesh, manual_axes, check_vma, spec, self) def unshard(self, mesh, check_vma, spec): return unshard_aval(mesh, check_vma, spec, self) def reshard(self, sharding): return reshard_aval(self, sharding) def vspace_add(self, x, y): from jax._src.ad_util import add_jaxvals # pyrefly: ignore[missing-import] return add_jaxvals(x, y) def raise_val2(self, lo_vals_ft): return self.raise_val(*lo_vals_ft.unflatten()) # pyrefly: ignore[missing-attribute] def lower_val2(self, hi_val): return ft.flatten(self.lower_val(hi_val)) # pyrefly: ignore[missing-attribute] InputType = tuple[AbstractValue, ...] OutputType = tuple[AbstractValue, ...] # For use in typing annotations to denote either a Tracer or a `valid_jaxtype`. Value = Any def valid_jaxtype(x) -> bool: try: aval = typeof(x) except TypeError: return False else: if hasattr(aval, "dtype") and aval.dtype == dtypes.string_dtype: return False else: return True def mem_kind_to_space(mem_kind: str | None) -> MemorySpace: if mem_kind == 'pinned_host': return MemorySpace.Host return MemorySpace.Device def mem_space_to_kind(mem_space: Any) -> str: """Converts a memory space to its corresponding XLA memory kind string. Supports standard MemorySpace enums and custom memory spaces that define `memory_kind` property. """ if isinstance(mem_space, MemorySpace): if mem_space == MemorySpace.Device: return "device" elif mem_space == MemorySpace.Host: return "pinned_host" elif hasattr(mem_space, "memory_kind"): return mem_space.memory_kind assert False, f"unreachable: {mem_space}" @cache(max_size=4096, trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value) def update_aval_with_sharding(aval, sharding, mat=None): if isinstance(sharding, NamedSharding): s = NamedSharding(sharding.mesh.abstract_mesh, sharding.spec._normalized_spec_for_aval(aval.ndim)) return aval.update( sharding=s, manual_axis_type=(aval.mat if mat is None else mat), memory_space=mem_kind_to_space(sharding.memory_kind)) return aval if mat is None else aval.update(manual_axis_type=mat) # We have two flavors of abstractification APIs here which each used to have # their own separate implementation. Now they're effectively the same, with the # following differences: # # - typeof returns avals for valid array-like objects, including tracers. # - shaped_abstractify is like typeof, but also accepts duck-typed arrays. # def shaped_abstractify(x): typ = type(x) if (aval_fn := pytype_aval_mappings.get(typ)): # fast path return aval_fn(x) for t in typ.__mro__[1:]: if (aval_fn := pytype_aval_mappings.get(t)): return aval_fn(x) if isinstance(x, AbstractValue): return x if getattr(x, '__jax_array__', None) is not None: raise ValueError( 'Triggering __jax_array__() during abstractification is no longer' ' supported. To avoid this error, either explicitly convert your object' ' using jax.numpy.array(), or register your object as a pytree.' ) if hasattr(x, 'dtype'): aval = ShapedArray( np.shape(x), dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True), weak_type=getattr(x, "weak_type", False), ) return update_aval_with_sharding(aval, getattr(x, 'sharding', None)) raise TypeError( f"Cannot interpret value of type {typ} as an abstract array; it " "does not have a dtype attribute") # TODO(phawkins): the return type should be AbstractValue. def typeof(x: Any) -> Any: """Return the JAX type (i.e. :class:`AbstractValue`) of the input. Raises a ``TypeError`` if ``x`` is not a valid JAX type. """ typ = type(x) if (aval_fn := pytype_aval_mappings.get(typ)): # fast path return aval_fn(x) for t in typ.__mro__[1:]: if (aval_fn := pytype_aval_mappings.get(t)): return aval_fn(x) if getattr(x, '__jax_array__', None) is not None: raise ValueError( 'Triggering __jax_array__() during abstractification is no longer' ' supported. To avoid this error, either explicitly convert your object' ' using jax.numpy.array(), or register your object as a pytree.' ) raise TypeError(f"Argument '{x}' of type '{typ}' is not a valid JAX type") def is_concrete(x): return to_concrete_value(x) is not None def to_concrete_value(x): if isinstance(x, Tracer): return x.to_concrete_value() else: return x def concretization_function_error(fun, suggest_astype=False): fname = getattr(fun, "__name__", fun) fname_context = f"The problem arose with the `{fname}` function. " if suggest_astype: fname_context += ("If trying to convert the data type of a value, " f"try using `x.astype({fun.__name__})` " f"or `jnp.array(x, {fun.__name__})` instead.") if fun is bool: def error(self, arg): raise TracerBoolConversionError(arg) elif fun in (hex, oct, operator.index): def error(self, arg): raise TracerIntegerConversionError(arg) else: def error(self, arg): raise ConcretizationTypeError(arg, fname_context) return error def concrete_or_error(force: Any, val: Any, context=""): """Like force(val), but gives the context in the error message.""" if force is None: force = lambda x: x if isinstance(val, Tracer): maybe_concrete = val.to_concrete_value() if maybe_concrete is None: raise ConcretizationTypeError(val, context) else: return force(maybe_concrete) else: return force(val) def concrete_dim_or_error(val: Any, context=""): """Like concrete_or_error(operator.index), allowing symbolic dimensions.""" if is_symbolic_dim(val): return val else: return concrete_or_error(operator.index, val, context=context) ### Extended dtypes # # Extended dtypes are JAX-specific dtypes that allow us to represent logical # arrays of element types that do not have an obvious direct correspondence # to ("physical") arrays of basic types in a compiler. In particular, their # element types differ from those of XLA and NumPy (e.g. int32). These dtypes # are only known to JAX. Their implementation is determined by: # a) an object representing the extended dtype, accessible via the `dtype` # attribute on corresponding JAX arrays and, internally, on avals such # as ShapedArrays that correspond to such JAX arrays; # b) a set of rules, available via a private attribute on the extended dtype # object in (a). # The rules in (b) tell JAX internals how to ground out the element # type for interaction with the compiler and runtime, e.g. when lowering # to the compiler's language. @overload def physical_aval(aval: ShapedArray) -> ShapedArray: ... @overload # TODO(frostig): remove this case def physical_aval(aval: AbstractValue) -> AbstractValue: ... def physical_aval(aval): if (isinstance(aval, ShapedArray) and isinstance(aval.dtype, dtypes.ExtendedDType)): elt_aval = physical_element_aval(aval.dtype) from jax._src.sharding_impls import physical_sharding # pyrefly: ignore[missing-import] return ShapedArray((*aval.shape, *elt_aval.shape), elt_aval.dtype, sharding=physical_sharding(aval, aval.sharding), manual_axis_type=aval.mat, memory_space=aval.memory_space) return aval def physical_shape(logical_shape, dtype): elt_aval = physical_element_aval(dtype) return (*logical_shape, *elt_aval.shape) def physical_element_aval(edtype: dtypes.ExtendedDType) -> ShapedArray: duck = edtype._rules.physical_element_aval(edtype) return ShapedArray(duck.shape, dtypes.dtype(duck.dtype)) _dtype_object_types = (np.dtype, dtypes.ExtendedDType) def _dtype_object(dtype): return dtype if isinstance(dtype, _dtype_object_types) else np.dtype(dtype) def _canonicalize_dimension(dim: DimSize) -> DimSize: # Dimensions are most commonly integral (by far), so we check that first. try: return operator.index(dim) except TypeError as e: type_error = e if is_dim(dim): return dim else: raise type_error def canonicalize_shape(shape: Shape, context: str="") -> tuple[Any, ...]: """Canonicalizes and checks for errors in a user-provided shape value. Args: shape: a Python value that represents a shape. Returns: A tuple of canonical dimension values. """ if isinstance(shape, int): shape = shape, try: return tuple(unsafe_map(_canonicalize_dimension, shape)) except TypeError: pass raise _invalid_shape_error(shape, context) def canonicalize_dim(d: DimSize, context: str="") -> DimSize: """Canonicalizes and checks for errors in a user-provided shape dimension value. Args: d: a Python value that represents a dimension. Returns: A canonical dimension value. """ return canonicalize_shape((d,), context)[0] def _invalid_shape_error(shape: Shape, context: str=""): msg = ("Shapes must be 1D sequences of concrete values of integer type, " f"got {shape}.") if context: msg += f" {context}." if any(isinstance(x, Tracer) and isinstance(typeof(x), ShapedArray) and not is_concrete(x) for x in shape): msg += ("\nIf using `jit`, try using `static_argnums` or applying `jit` to " "smaller subfunctions.") for x in shape: if isinstance(x, Tracer) and hasattr(x, "_origin_msg"): msg += x._origin_msg() return TypeError(msg) class ShardingTypeError(Exception): pass class MemorySpace(enum.Enum): Device = enum.auto() Host = enum.auto() Any = enum.auto() def __repr__(self): return f"MemorySpace.{self.name}" # For reasons passing my understanding Enum.__hash__ hashes the name string. # Use an object identity hash instead. __hash__ = object.__hash__ def get_cur_mesh_sharding(spec=None): spec = P() if spec is None else spec return NamedSharding(mesh_lib.get_abstract_mesh(), spec) def getu(aval, kind=UnreducedKind.sum): if aval.sharding.mesh.are_all_axes_manual: out_u = aval.mat.unreduced if out_u: if (aval_k := aval.mat.unreduced_kind) is not kind: raise ValueError(f'Expected unreduced_kind={kind} but got {aval_k}') return out_u if aval.sharding.mesh.are_all_axes_explicit: out_u = aval.sharding.spec.unreduced if out_u: if (aval_k := aval.sharding.spec.unreduced_kind) is not kind: raise ValueError(f'Expected unreduced_kind={kind} but got {aval_k}') return out_u # Revise this after partial manual unreduced is supported assert not aval.mat.unreduced return aval.sharding.spec.unreduced def getr(aval): if aval.sharding.mesh.are_all_axes_manual: return aval.mat.reduced if aval.sharding.mesh.are_all_axes_explicit: return aval.sharding.spec.reduced # Revise this after partial manual reduced is supported assert not aval.mat.reduced return aval.sharding.spec.reduced def _make_lengths_same(sharding, ndim): pspec = sharding.spec if ndim > len(pspec): return sharding.update(spec=pspec._normalized_spec_for_aval(ndim)) if ((sharding.mesh.empty or sharding.mesh._are_all_axes_auto_or_manual) and ndim < len(pspec)): assert all(s is None for s in pspec.partitions[ndim:]) return sharding.update(spec=sharding.spec.update( partitions=pspec.partitions[:ndim])) return sharding def modify_spec_for_auto_manual(spec, mesh) -> P: new_spec: list[Any] = [] # PartitionSpec can only mention mesh axes that are Explicit. for s in spec.partitions: if s is None: new_spec.append(s) elif isinstance(s, tuple): new_spec.append(tuple( p for p in s if mesh._name_to_type[p] == AxisType.Explicit)) else: new_spec.append(s if mesh._name_to_type[s] == AxisType.Explicit else None) new_unreduced = {u for u in spec.unreduced if mesh._name_to_type[u] == AxisType.Explicit} new_reduced = {u for u in spec.reduced if mesh._name_to_type[u] == AxisType.Explicit} u_kind = spec.unreduced_kind if new_unreduced else None return P(*new_spec, unreduced=new_unreduced, reduced=new_reduced, unreduced_kind=u_kind) def _maybe_modify_sharding(sharding, ndim): if len(sharding.spec) == 0 or all(s is None for s in sharding.spec.partitions): out = sharding elif sharding.mesh.are_all_axes_explicit: out = sharding else: out = sharding.update(spec=modify_spec_for_auto_manual( sharding.spec, sharding.mesh)) if config.remove_size_one_mesh_axis_from_type.value: out = out.update(spec=ns.remove_size_one_mesh_axis_from_spec(out.spec, out.mesh)) if len(out.spec) != ndim: out = _make_lengths_same(out, ndim) return out def _check_divisibility(sharding, shape): mesh = sharding.mesh for dim, (spec, sh) in enumerate(zip(sharding.spec.partitions, shape)): if spec is None: continue spec = spec if isinstance(spec, tuple) else (spec,) size = math.prod(mesh.shape[s] for s in spec) _, remainder = divmod(sh, size) if remainder != 0: raise ValueError( f"Sharding spec {spec} implies that array axis {dim} is partitioned" f" {size} times, but does not evenly divide the dimension size {sh}." f" Got shape: {shape} and sharding {sharding}") @cache(max_size=4096, trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value) def get_sharding(sharding, shape): """Modifies and checks the sharding. Some modifications/checks include: * Making the length of specs the same as ndim * If a mesh axis is mentioned in pspec is Auto/Manual, replace it with None * Checking for len(spec)-ndim match * Checking if the mesh is an AbstractMesh. """ ndim = len(shape) if sharding is None: return _empty_sharding(ndim) out_s = _maybe_modify_sharding(sharding, ndim) if len(out_s.spec) != ndim: raise ValueError( f"Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's" f" ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and" f" sharding {out_s}") if not isinstance(out_s.mesh, mesh_lib.AbstractMesh): raise ValueError("Mesh of an aval must be an AbstractMesh. " f"Got {out_s.mesh} of type {type(out_s.mesh)}") _check_divisibility(out_s, shape) if out_s.memory_kind is not None: raise ValueError( "sharding with memory_kind is not allowed. Please use `jax.device_put`" f" to transfer to different memory spaces. Got {sharding=}") return out_s @cache(max_size=4096, trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value) def get_mat(mat, mesh): if mesh.empty: assert mat.empty, mat return mat axis_env = get_axis_env() in_axis_env = lambda i: axis_env.axis_exists(i) and i not in mesh._name_to_type for i in it.chain(mat.varying, mat.unreduced, mat.reduced): if in_axis_env(i): continue if mesh._name_to_type[i] != AxisType.Manual: raise ValueError( "Axes mentioned in `manual_axis_type` field of ShapedArray should be" f" of type `Manual`. Got manual_axis_type={mat} with axis: {i} of" f" type {mesh._name_to_type[i]}") if config.remove_size_one_mesh_axis_from_type.value: varying = frozenset(i for i in mat.varying if in_axis_env(i) or mesh.shape[i] != 1) unreduced = frozenset(u for u in mat.unreduced if mesh.shape[u] != 1) reduced = frozenset(r for r in mat.reduced if mesh.shape[r] != 1) u_kind = mat.unreduced_kind if unreduced else None return mat.update(varying=varying, unreduced=unreduced, reduced=reduced, unreduced_kind=u_kind) return mat def get_memory_space(memory_space): assert memory_space is not None return memory_space def _check_mat(varying, unreduced, reduced, unreduced_kind): if varying & unreduced: raise ValueError( "varying and unreduced cannot have common mesh axes. Got" f" varying={varying} and unreduced={unreduced}") if varying & reduced: raise ValueError( "varying and reduced cannot have common mesh axes. Got" f" varying={varying} and reduced={reduced}") assert not (varying & unreduced & reduced) if unreduced_kind is not None and not isinstance(unreduced_kind, UnreducedKind): raise TypeError( "Expected unreduced_kind to be of type `jax.sharding.UnreducedKind`" f" but got {type(unreduced_kind)}") if not unreduced and unreduced_kind is not None: raise ValueError( "`unreduced_kind` should be `None` when `unreduced` is an empty set." f" Got {unreduced_kind=} and {unreduced=}") def _canonicalize_mat(name, val): if not isinstance(val, frozenset): if not isinstance(val, set): raise TypeError( f"{name} argument of ManualAxisType should " f"of type `frozenset` or `set`. Got type {type(val)}") val = frozenset(val) return val @immutable class ManualAxisType: __slots__ = ('varying', 'unreduced', 'reduced', 'unreduced_kind', '__weakref__') varying: frozenset unreduced: frozenset reduced: frozenset unreduced_kind: UnreducedKind | None @staticmethod @weak_value_interner def _create(varying, unreduced, reduced, unreduced_kind): # We cannot modify the arguments within the interned function, but we are # free to throw an exception. _check_mat(varying, unreduced, reduced, unreduced_kind) obj = object.__new__(ManualAxisType) object.__setattr__(obj, 'varying', varying) object.__setattr__(obj, 'unreduced', unreduced) object.__setattr__(obj, 'reduced', reduced) object.__setattr__(obj, 'unreduced_kind', unreduced_kind) return obj def __new__(cls, *, varying=frozenset(), unreduced=frozenset(), reduced=frozenset(), unreduced_kind: UnreducedKind | None = None): varying = _canonicalize_mat('varying', varying) unreduced = _canonicalize_mat('unreduced', unreduced) reduced = _canonicalize_mat('reduced', reduced) if unreduced and unreduced_kind is None: unreduced_kind = UnreducedKind.sum return cls._create(varying, unreduced, reduced, unreduced_kind) # No __eq__ or __hash__: interned classes use object identity. def __repr__(self): return (f"ManualAxisType(varying={self.varying}, " f"unreduced={self.unreduced}, reduced={self.reduced}, " f"unreduced_kind={self.unreduced_kind})") def __getnewargs_ex__(self): return (), {'varying': self.varying, 'unreduced': self.unreduced, 'reduced': self.reduced, 'unreduced_kind': self.unreduced_kind} def update(self, **kwargs): if 'varying' not in kwargs: kwargs['varying'] = self.varying if 'unreduced' not in kwargs: kwargs['unreduced'] = self.unreduced if 'reduced' not in kwargs: kwargs['reduced'] = self.reduced if 'unreduced_kind' not in kwargs: kwargs['unreduced_kind'] = self.unreduced_kind return ManualAxisType(**kwargs) def to_ct_mat(self): assert self.unreduced_kind is None or self.unreduced_kind is UnreducedKind.sum kind = UnreducedKind.sum if self.reduced else None return self.update(unreduced=self.reduced, reduced=self.unreduced, unreduced_kind=kind) @property def empty(self): return self is empty_mat def invarying(self, mesh) -> frozenset: return frozenset(mesh.manual_axes) - ( self.varying | self.unreduced | self.reduced) @property def vur(self) -> frozenset: return self.varying | self.unreduced | self.reduced empty_mat = ManualAxisType() @functools.cache def _empty_sharding(ndim): return NamedSharding(mesh_lib.empty_abstract_mesh, P(*[None] * ndim)) @immutable class ShapedArray(AbstractValue): # inherits slots from parent __slots__ = ['shape', 'dtype', 'weak_type', 'sharding', 'manual_axis_type', 'memory_space', 'layout', '_stripped_weak_type', '__weakref__'] array_abstraction_level = 2 shape: Any dtype: Any weak_type: Any sharding: Any manual_axis_type: Any memory_space: Any layout: Any _stripped_weak_type: Any @staticmethod @weak_value_interner def _create(shape, dtype, weak_type, sharding, manual_axis_type, memory_space, layout): obj = object.__new__(ShapedArray) object.__setattr__(obj, 'shape', shape) object.__setattr__(obj, 'dtype', dtype) object.__setattr__(obj, 'weak_type', weak_type) object.__setattr__(obj, 'sharding', sharding) object.__setattr__(obj, 'manual_axis_type', manual_axis_type) object.__setattr__(obj, 'memory_space', memory_space) object.__setattr__(obj, 'layout', layout) object.__setattr__(obj, '_stripped_weak_type', None) return obj def __new__(cls, shape, dtype, weak_type=False, *, sharding=None, manual_axis_type: ManualAxisType = empty_mat, memory_space: MemorySpace = MemorySpace.Device, layout=AutoLayout): shape = canonicalize_shape(shape) dtype = _dtype_object(dtype) if sharding is None: sharding = _empty_sharding(len(shape)) assert manual_axis_type.empty, manual_axis_type else: sharding = get_sharding(sharding, shape) # https://docs.jax.dev/en/latest/201/shard-map.html#tracking-how-values-vary-over-manual-mesh-axes-and-check-vma-true manual_axis_type = get_mat(manual_axis_type, sharding.mesh) # See description of https://github.com/jax-ml/jax/pull/30556 memory_space = get_memory_space(memory_space) return cls._create(shape, dtype, weak_type, sharding, manual_axis_type, memory_space, layout) # Interned types don't need __eq__ or __hash__. @property def mat(self): return self.manual_axis_type def update(self, shape=None, dtype=None, weak_type=None, **kwargs): if shape is None: shape = self.shape if dtype is None: dtype = self.dtype if weak_type is None: weak_type = self.weak_type if 'sharding' not in kwargs: kwargs['sharding'] = self.sharding if 'manual_axis_type' not in kwargs: kwargs['manual_axis_type'] = self.manual_axis_type if 'memory_space' not in kwargs: kwargs['memory_space'] = self.memory_space if 'layout' not in kwargs: kwargs['layout'] = self.layout return ShapedArray(shape, dtype, weak_type, **kwargs) ndim = property(lambda self: len(self.shape)) size = property(lambda self: 0 if any(type(d) is int and d == 0 for d in self.shape) else math.prod(self.shape)) broadcast: ClassVar[aval_method | None] = None transpose: ClassVar[aval_method | None] = None reshape: ClassVar[aval_method | None] = None _iter: ClassVar[staticmethod | None] = None def __getnewargs_ex__(self): return (self.shape, self.dtype, self.weak_type), { 'sharding': self.sharding, 'manual_axis_type': self.manual_axis_type, 'memory_space': self.memory_space, 'layout': self.layout, } def str_short(self, short_dtypes=False, mesh_axis_types=False): return str_short_aval( self.shape, self.dtype, self.sharding.mesh, self.sharding.spec, self.mat, self.memory_space, short_dtypes, mesh_axis_types) def __repr__(self): wt_str = ", weak_type=True" if self.weak_type else "" return f'ShapedArray({self.str_short()}{wt_str})' def __str__(self): wt_str = "~" if self.weak_type else "" return f'{wt_str}{self.str_short()}' def to_tangent_aval(self): return ShapedArray._create( self.shape, primal_dtype_to_tangent_dtype(self.dtype), self.weak_type, self.sharding, self.mat, self.memory_space, self.layout) def to_ct_aval(self): dtype = primal_dtype_to_tangent_dtype(self.dtype) sharding = primal_sharding_to_cotangent_sharding(self.sharding) ct_mat = self.mat.to_ct_mat() return ShapedArray._create( self.shape, dtype, self.weak_type, sharding, ct_mat, self.memory_space, self.layout) def _len(self, ignored_tracer): try: return self.shape[0] except IndexError as err: raise TypeError("len() of unsized object") from err # same as numpy error def update_manual_axis_type(self, mat): mat = get_mat(mat, self.sharding.mesh) if mat is self.manual_axis_type: return self return ShapedArray._create(self.shape, self.dtype, self.weak_type, self.sharding, mat, self.memory_space, self.layout) def update_weak_type(self, weak_type): if weak_type == self.weak_type: return self return ShapedArray._create(self.shape, self.dtype, weak_type, self.sharding, self.manual_axis_type, self.memory_space, self.layout) def strip_weak_type(self) -> AbstractValue: if not self.weak_type: return self # _stripped_weak_type is not protected by a lock, but the access should be # safe because ShapedArray values are interned: if two threads race to set # the value it will be the same object. val = self._stripped_weak_type if val is None: val = self.update_weak_type(False) object.__setattr__(self, '_stripped_weak_type', val) return val def nospec(self, mesh, check_vma, all_names) -> P: # TODO(mattjj, yashkatariya): should use newly all_names in check_vma path? sh_names = (order_wrt_mesh(mesh, self.mat.varying) if check_vma else all_names) u_names = self.mat.unreduced if check_vma else frozenset() r_names = self.mat.reduced if check_vma else frozenset() u_kind = self.mat.unreduced_kind if check_vma else None return (P(sh_names, unreduced=u_names, reduced=r_names, unreduced_kind=u_kind) if sh_names else P(unreduced=u_names, reduced=r_names, unreduced_kind=u_kind)) _bool = concretization_function_error(bool) _int = concretization_function_error(int, True) _float = concretization_function_error(float, True) _complex = concretization_function_error(complex, True) _hex = concretization_function_error(hex) _oct = concretization_function_error(oct) _index = concretization_function_error(operator.index) def _get_shape_sharding_str(shape, spec): out = [] for s1, s2 in zip(shape, spec.partitions): if s2 is None: out.append(f"{s1}") elif isinstance(s2, tuple): ss = ','.join(s for s in s2) out.append(f"{s1}@({ss})") else: out.append(f"{s1}@{s2}") return ','.join(out) @cache(max_size=1024, trace_context_in_key=False) def _axis_types_dict(mesh): if not mesh.axis_names: return {} d = defaultdict(list) for n, t in safe_zip(mesh.axis_names, mesh.axis_types): d[t].append(n) return {t: tuple(n) for t, n in d.items()} def str_short_aval(shape, dtype, mesh, spec, mat, memory_space, short_dtypes=False, mesh_axis_types=False) -> str: dt_str = dtypes.short_dtype_name(dtype) if short_dtypes else dtype.name dt_str = dt_str.replace('void', 'float0') shapestr = _get_shape_sharding_str(shape, spec) mesh_axes = f'({_axis_types_dict(mesh)})' if mesh_axis_types else '' vma_ur = _vma_ur_str(mat, spec.unreduced, spec.reduced, spec.unreduced_kind, mesh) ms_str = ("" if memory_space == MemorySpace.Device else f"<{memory_space.name.lower()}>") return f'{dt_str}{ms_str}[{shapestr}]{vma_ur}{mesh_axes}' def _create_str(x, prefix): x_str = f"{','.join(i for i in x)}" x_str = x_str if len(x) == 1 else f"({x_str})" return f"{prefix}:{x_str}, " def order_wrt_mesh(mesh, x): return tuple(a for a in mesh.axis_names if a in x) def _vma_ur_str(mat, spec_unreduced, spec_reduced, u_kind, mesh): vma = mat.varying # TODO(yashkatariya): Diff between explicit unreduced and manual unreduced unreduced = mat.unreduced | spec_unreduced reduced = mat.reduced | spec_reduced if not vma and not unreduced and not reduced: return '' vma_str = _create_str(order_wrt_mesh(mesh, vma), 'V') if vma else '' u_str = '' if unreduced: u_prefix = ('U' if u_kind is None or u_kind is UnreducedKind.sum else f'U_{u_kind.name}') u_str = _create_str(order_wrt_mesh(mesh, unreduced), u_prefix) r_str = _create_str(order_wrt_mesh(mesh, reduced), 'R') if reduced else '' m_str = f"{vma_str}{u_str}{r_str}".rstrip(', ') return f"{{{m_str}}}" def primal_dtype_to_tangent_dtype(primal_dtype): if isinstance(primal_dtype, dtypes.ExtendedDType): return primal_dtype._rules.tangent_dtype(primal_dtype) elif not dtypes.issubdtype(primal_dtype, np.inexact): return dtypes.float0 else: return primal_dtype def primal_sharding_to_cotangent_sharding(sharding): return sharding.update(spec=sharding.spec.to_ct_spec()) ############################## pvary ################################# # Invariant -> Variant no-op cast def pvary(x, axis_name): axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name if not axis_name: return x cur_mesh = mesh_lib.get_abstract_mesh() if not config._check_vma.value and all(a in cur_mesh.manual_axes for a in axes): return x new_axes = axes if cur_mesh.empty else order_wrt_mesh(cur_mesh, axes) assert set(new_axes) == set(axes) del axes # TODO(yashkatariya): Remove this handling and remove_size_one_mesh_axis_from_type # generally from JAX. if config.remove_size_one_mesh_axis_from_type.value and not cur_mesh.empty: new_axes = tuple(i for i in new_axes if cur_mesh.shape[i] != 1) if not new_axes: return x return tree_map(lambda leaf: pvary_p.bind(leaf, axes=new_axes), x) pvary_p = Primitive('pvary') ####################### reduced_vary_cast ############################# # Reduced -> Varying no-op cast def reduced_vary_cast(x, axis_name): axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name if not axis_name: return x cur_mesh = mesh_lib.get_abstract_mesh() if not config._check_vma.value and all(a in cur_mesh.manual_axes for a in axes): return x new_axes = axes if cur_mesh.empty else order_wrt_mesh(cur_mesh, axes) assert set(new_axes) == set(axes) del axes return tree_map(lambda leaf: reduced_vary_cast_p.bind(leaf, axes=new_axes), x) reduced_vary_cast_p = Primitive('reduced_vary_cast') ####################################################################### def check_unreduced_args(args, axes, name, kind=UnreducedKind.sum): axes = axes if isinstance(axes, (tuple, list)) else (axes,) axes = set(axes) for a in args: if a.mat.unreduced & axes: raise ValueError( f"{name} cannot accept args which are unreduced. Got" f" {a.str_short(True)} and axes={axes}") if a.mat.unreduced and a.mat.unreduced_kind is not kind: raise ValueError( f"{name} cannot accept args with" f" unreduced_kind={a.mat.unreduced_kind}. Expected" f" unreduced_kind={kind}") if a.mat.reduced & axes: raise ValueError( f"{name} cannot accept args which are reduced. Got" f" {a.str_short(True)} and axes={axes}") def insert_reduced_reshard(args): cur_mesh = mesh_lib.get_abstract_mesh() if not cur_mesh.are_all_axes_explicit: return args # TODO(yashkatariya): Handle >2 args too if len(args) != 2: return args in_reduced = [aval.sharding.spec.reduced if isinstance(aval := shaped_abstractify(a), ShapedArray) else frozenset() for a in args] out_reduced = frozenset.union(*in_reduced) out = [] for arg, src_reduced in zip(args, in_reduced): aval = shaped_abstractify(arg) if (isinstance(aval, ShapedArray) and aval.ndim == 0 and out_reduced and (get_replicated_axes(aval.sharding.spec, cur_mesh) & out_reduced) == out_reduced): from jax._src.pjit import reshard # type: ignore out.append(reshard(arg, P(reduced=out_reduced))) else: out.append(arg) return out def auto_insert_reshard(*args): if not args: return args if not config._check_vma.value: return insert_reduced_reshard(args) in_vma = [aval.mat.varying if isinstance(aval := typeof(a), ShapedArray) else frozenset() for a in args] in_reduced = [aval.mat.reduced if isinstance(aval := typeof(a), ShapedArray) else frozenset() for a in args] out_vma = frozenset.union(*in_vma) out = [] for arg, src_vma, src_reduced in zip(args, in_vma, in_reduced): if (isinstance(aval := typeof(arg), ShapedArray) and (rest_vma := out_vma - src_vma)): # TODO(yashkatariya): Handle partial reduced_vary_cast and partial pvary. # Will need more changes to pvary to allow such partialness. if not config.auto_pcast.value and dtypes.issubdtype(aval.dtype, np.inexact): out.append(arg) elif src_reduced == rest_vma: out.append( reduced_vary_cast(arg, tuple(n for n in out_vma if n in rest_vma))) else: out.append(pvary(arg, tuple(n for n in out_vma if n in rest_vma))) else: out.append(arg) return out def standard_vma_rule(prim_name, *avals, **kwargs) -> frozenset[AxisName]: if not config._check_vma.value: return frozenset() avals = tuple(a for a in avals if a is not abstract_token) if not avals: return frozenset() vma, *vmas = (a.mat.varying for a in avals) if not all(vma == vma_ for vma_ in vmas): raise ValueError( f'Primitive {prim_name} requires varying manual axes ' f'to match, but got {[vma, *vmas]}. Please open an issue at ' 'https://github.com/jax-ml/jax/issues and as a temporary ' 'workaround pass the check_vma=False argument to `jax.shard_map`') return vma @dataclass(frozen=True, slots=True) class bint(dtypes.ExtendedDType): bound: int @property def type(self) -> type: return dtypes.extended @property def name(self) -> str: return f'bint{{≤{self.bound}}}' def __str__(self) -> str: return self.name AxisSize = int | Tracer | Var class RefMeta(type): def __instancecheck__(self, inst): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] return (super().__instancecheck__(inst) or isinstance(inst, Tracer) and isinstance(inst.aval, AbstractRef)) class Ref(metaclass=RefMeta): """Mutable array reference. In most cases this should not be constructed directly, but rather via :func:`jax.ref.new_ref`. For examples of how this can be used, refer to the `Ref guide`_. .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs """ _aval: AbstractValue _refs: PyTree # list of ArrayRefImpl def __init__(self, aval, refs): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] assert isinstance(aval, AbstractRef) self._aval = aval self._refs = refs def __repr__(self) -> str: if self._aval.is_high: return f"Ref({self._refs})" return "Ref" + repr(self._refs._buf)[5:] # forward type-level info to aval aval = property(lambda self: self._aval) shape = property(lambda self: self._aval.shape) size = property(lambda self: self._aval.size) ndim = property(lambda self: len(self._aval.shape)) dtype = property(lambda self: self._aval.dtype) # get operations from aval, munging the name def __getitem__(self, idx): return self._aval._getitem(self, idx) # pyrefly: ignore[missing-attribute] def __setitem__(self, idx, x): return self._aval._setitem(self, idx, x) # pyrefly: ignore[missing-attribute] def __len__(self) -> int: return self._aval._len(self) # pyrefly: ignore[missing-attribute] def addupdate(self, x, idx=()): return self._aval._addupdate(self, idx, x) # pyrefly: ignore[missing-attribute] # some attributes/methods only work for lojax refs sharding = property(lambda self: self._refs._buf.sharding) format = property(lambda self: self._refs._buf.format) committed = _committed = property(lambda self: True) def unsafe_buffer_pointer(self): return self._refs._buf.unsafe_buffer_pointer() @property def at(self): raise NotImplementedError() # TODO(mattjj) class ArrayRefImpl: _aval: AbstractValue _buf: Array # mutable field def __init__(self, aval, buf): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] assert isinstance(aval, AbstractRef) and isinstance(aval.inner_aval, ShapedArray) self._aval = aval self._buf = buf pytype_aval_mappings[Ref] = lambda x: x._aval dtypes.register_canonicalize_value_handler(Ref, None) class InternalMutableArrayEffect(effects.Effect): pass array_ref_effect = internal_mutable_array_effect = InternalMutableArrayEffect() effects.control_flow_allowed_effects.add_type(InternalMutableArrayEffect) effects.remat_allowed_effects.add_type(InternalMutableArrayEffect) def new_ref(init_val: Any, *, memory_space: Any = None, kind: Any = None, pin: bool = False): """Create a mutable array reference with initial value ``init_val``. For more discussion, see the `Ref guide`_. Args: init_val: A :class:`jax.Array` representing the initial state of the buffer. memory_space: An optional memory space attribute for the Ref. kind: An optional string indicating the mutation semantics under rematerialization. pin: Whether to lower the ref to a pinned buffer in HLO. Returns: A :class:`jax.ref.Ref` containing a reference to a mutable buffer. .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs """ return ref_p.bind(init_val, memory_space=memory_space, kind=kind, pin=pin) ref_p = Primitive('new_ref') ref_p.is_effectful = lambda params: True ref_p.ref_primitive = True ref_p.ref_allocating = True ref_p.is_high = lambda aval, *, memory_space, kind, pin: aval.is_high def _ref_to_lojax(init_val, *, memory_space, kind, pin): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] val_ty = typeof(init_val) hival_of_refs = val_ty.raise_val(*map(new_ref, val_ty.lower_val(init_val))) return Ref(AbstractRef(val_ty), hival_of_refs) ref_p.to_lojax = _ref_to_lojax @ref_p.def_effectful_abstract_eval def _ref_abstract_eval(init_aval, *, memory_space: Any, kind: Any, pin: bool): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] # If no memory space is specified, use the memory space of the initial value # but we make sure to reset it to Device because the Ref owns the memory space if isinstance(init_aval, ShapedArray): if ( memory_space is None and init_aval.memory_space is not MemorySpace.Device ): memory_space = init_aval.memory_space init_aval = init_aval.update( memory_space=MemorySpace.Device, weak_type=False ) return (AbstractRef(init_aval, memory_space=memory_space, kind=kind), {internal_mutable_array_effect}) @ref_p.def_impl def _ref_impl(init_val, *, memory_space: Any, kind: Any, pin: bool): if memory_space not in (None, MemorySpace.Device, MemorySpace.Host): raise NotImplementedError( "array ref with memory space only works inside of a `jit`.") if pin: raise NotImplementedError( "pinned array ref only works inside of a `jit`.") from jax._src.api import device_put # pyrefly: ignore[missing-import] from jax._src.lax.lax import _array_copy, _convert_element_type # pyrefly: ignore[missing-import] if memory_space is not None: init_val = device_put(init_val, memory_space) if dtypes.is_weakly_typed(init_val): buf = _convert_element_type( init_val, dtypes.dtype(init_val), weak_type=False ) else: buf = _array_copy(init_val) aval, _ = _ref_abstract_eval( typeof(buf), memory_space=memory_space, kind=kind, pin=pin ) if not buf.committed: buf = device_put(buf, buf.sharding) return Ref(aval, ArrayRefImpl(aval, buf)) # TODO(mattjj,dougalm): merge with ref_p def empty_ref(ty, memory_space=None, pin=False): aval = shaped_abstractify(ty) return empty_ref_p.bind(ty=aval, memory_space=memory_space, pin=pin) empty_ref_p = Primitive('empty_ref') empty_ref_p.ref_primitive = True empty_ref_p.is_effectful = lambda _: True empty_ref_p.ref_allocating = True empty_ref_p.is_high = lambda *, ty, memory_space, pin: ty.is_high def _empty_ref_to_lojax(*, ty, memory_space, pin): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] n = len(ty.lo_ty()) hival_of_refs = ty.raise_val( *map(empty_ref, ty.lo_ty(), [memory_space] * n, [pin] * n)) return Ref(AbstractRef(ty), hival_of_refs) empty_ref_p.to_lojax = _empty_ref_to_lojax @empty_ref_p.def_effectful_abstract_eval def _empty_ref_abstract_eval(*, ty, memory_space, pin): from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] if isinstance(ty, ShapedArray) and ty.weak_type: ty = ty.update(weak_type=False) return (AbstractRef(ty, memory_space=memory_space), {internal_mutable_array_effect}) # TODO(mattjj,dougalm): merge with freeze_p def free_ref(ref: Ref): """Invalidate a given reference.""" free_ref_p.bind(ref) return () free_ref_p = Primitive('free_ref') free_ref_p.multiple_results = True free_ref_p.is_effectful = lambda _: True free_ref_p.ref_primitive = True @free_ref_p.def_effectful_abstract_eval def _free_ref_abstract_eval(ref_aval): # No effects, but there is a custom DCE rule that prevents free_ref from # being DCE'd. return (), {} @free_ref_p.def_impl def _free_ref_impl(ref): return () def freeze(ref: Ref) -> Array: """Invalidate a given reference and return its final value. For more information about mutable array references, refer to the `Ref guide`_. Args: ref: A :class:`jax.ref.Ref` object. Returns: A :class:`jax.Array` containing the contents of ``ref``. Examples: >>> import jax >>> ref = jax.new_ref(jax.numpy.arange(5)) >>> ref[3] = 100 >>> ref Ref([ 0, 1, 2, 100, 4], dtype=int32) >>> jax.ref.freeze(ref) Array([ 0, 1, 2, 100, 4], dtype=int32) .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs """ return freeze_p.bind(ref) freeze_p = Primitive('freeze') freeze_p.is_effectful = lambda params: True freeze_p.ref_primitive = True freeze_p.is_high = lambda aval: aval.is_high def _freeze_to_lojax(ref): aval = typeof(ref._refs) lovals = aval.lower_val(ref._refs) vals = [freeze(loval) for loval in lovals] return aval.raise_val(*vals) freeze_p.to_lojax = _freeze_to_lojax @freeze_p.def_effectful_abstract_eval def freeze_abstract_eval(ref_aval): return ref_aval.inner_aval, {internal_mutable_array_effect} @freeze_p.def_impl def _freeze_impl(ref): return ref._refs._buf def accum_grad_in_ref(x): return accum_grad_in_ref_p.bind(x) accum_grad_in_ref_p = Primitive('accum_grad_in_ref') accum_grad_in_ref_p.is_high = lambda *_: True accum_grad_in_ref_p.to_lojax = lambda x: x accum_grad_in_ref_p.def_abstract_eval(lambda x: x) accum_grad_in_ref_p.def_impl(lambda x: x) class AbstractToken(AbstractValue): def str_short(self, short_dtypes=False, mesh_axis_types=False): return 'Tok' def to_tangent_aval(self): return self def to_ct_aval(self): return self abstract_token: AbstractToken = AbstractToken() # Singleton shaped array used by all abstract tokens when shape/dtype is needed. def get_token_aval(): return ShapedArray((0,), np.dtype(np.bool_), sharding=None) # Concrete token object class Token: # The underlying data wrapped by the token, could be used to threaded in and # out of computations to build up data dependency. _buf: Array def __init__(self, buf): self._buf = buf def block_until_ready(self): self._buf.block_until_ready() pytype_aval_mappings[Token] = lambda _: abstract_token dtypes.register_canonicalize_value_handler(Token, None) class AbstractFuture(AbstractValue): def __init__(self, inner_aval, done_fun): self.inner_aval = inner_aval self.done_fun = done_fun def __eq__(self, other): return (isinstance(other, AbstractFuture) and self.inner_aval == other.inner_aval) def __hash__(self): return hash(self.inner_aval) def str_short(self, short_dtypes=False, mesh_axis_types=False) -> str: return f'AbstractFuture{{{self.inner_aval.str_short(True)}}}' ndim = property(lambda self: len(self.shape)) size = property(lambda self: math.prod(self.shape)) def update(self, **kwargs): return AbstractFuture(kwargs.pop("inner_aval", self.inner_aval), kwargs.pop("done_fun", self.done_fun)) @aval_method def done(tracer): return tracer.aval.done_fun(tracer) # type: ignore @property def shape(self): try: return self.inner_aval.shape except AttributeError: raise AttributeError(f"{self!r} has no `sharding`.") from None @property def dtype(self): try: return self.inner_aval.dtype except AttributeError: raise AttributeError(f"{self!r} has no `sharding`.") from None @property def sharding(self): try: return self.inner_aval.sharding except AttributeError: raise AttributeError(f"{self!r} has no `sharding`.") from None @property def manual_axis_type(self): try: return self.inner_aval.manual_axis_type except AttributeError: raise AttributeError(f"{self!r} has no `manual_axis_type`.") from None @property def mat(self): return self.manual_axis_type ### Operations on shapes and dimension sizes. @set_module("jax.errors") class InconclusiveDimensionOperation(Exception): """Raised when we cannot conclusively compute with symbolic dimensions.""" def is_symbolic_dim(v: Any) -> bool: """Checks if a value is a symbolic dimension used for shape polymorphism. This should be used very rarely, because symbolic dimensions overload all operators, and should just work. """ return getattr(v, "dimension_as_value", None) is not None def is_constant_dim(d: DimSize) -> bool: # Whether the dimension is a static integer constant. # Try using a fast path for non-concrete Tracers. if isinstance(d, Tracer) and not is_concrete(d): return False try: operator.index(d) return True except: return False def is_dim(v: Any) -> bool: return is_symbolic_dim(v) or is_constant_dim(v) def is_constant_shape(s: Shape) -> bool: # Whether the shape is a static constant. return all(is_constant_dim(d) for d in s) def definitely_equal_one_of_dim(d1: DimSize, dlist: Sequence[DimSize]) -> bool: return any(definitely_equal(d1, d) for d in dlist) def definitely_equal_shape(s1: Shape, s2: Shape) -> bool: """Check that two shapes are guaranteed to be element-wise equal. In presence of dynamic shapes may return False even when the shapes may be equal at runtime. """ return (len(s1) == len(s2) and all(unsafe_map(definitely_equal, s1, s2))) def divide_shape_sizes(s1: Shape, s2: Shape) -> DimSize: """Returns an integer "i" s.t., i * size(s2) == size(s1). Raises InconclusiveDimensionOperation if there is no such integer.""" sz1 = math.prod(s1) sz2 = math.prod(s2) if definitely_equal(sz1, sz2): # Takes care of sz1 and sz2 being 0 return 1 q, r = divmod(sz1, sz2) if isinstance(r, Tracer) or r != 0: raise InconclusiveDimensionOperation( f"Cannot divide evenly the sizes of shapes {tuple(s1)} and {tuple(s2)}. " f"The remainder {r} should be 0.") return q def cancel_divide_tracers(num, denom): partition = lambda l: partition_list([isinstance(d, Tracer) for d in l], l) num, num_tracers = partition(num) denom, denom_tracers = partition(denom) if num_tracers or denom_tracers: factor = _cancel_divide(num_tracers, denom_tracers) if factor is not None: size1 = math.prod(num) size2 = math.prod(denom) if size1 == size2 or size2 != 0: return factor * (size1 // size2 if size1 != size2 else 1) def _cancel_divide(num, denom): num = list(num) for a in denom: i = next((i for i, b in enumerate(num) if definitely_equal(a, b)), None) if i is None: break # couldn't cancel del num[i] else: return math.prod(num) def is_empty_shape(s: Shape) -> bool: return any(definitely_equal(d, 0) for d in s) def dilate_dim(d: DimSize, dilation: DimSize) -> DimSize: """max(0, 1 + dilation * (d - 1)). Assumes dilation >= 1. """ if definitely_equal(dilation, 1): # fast path return d return max_dim(1 + dilation * (d - 1), 0) def stride_dim(d: DimSize, window_size: DimSize, window_stride: DimSize) -> DimSize: """max(0, (d - window_size) // window_stride + 1) If d < window_size, returns 0. We assume window_size >= 1 and window_stride >= 1. """ # If d < window_size then (d - window_size) // window_stride < 0 return max_dim((d - window_size) // window_stride + 1, 0) def min_dim(d1: DimSize, d2: DimSize) -> DimSize: """Like min(d1, d2) but for both constant and symbolic dimensions.""" d1_is_constant = is_constant_dim(d1) if d1_is_constant and is_constant_dim(d2): return min(d1, d2) d1 = concrete_dim_or_error(d1, "argument `d1` of `core.min_dim`") d2 = concrete_dim_or_error(d2, "argument `d2` of `core.min_dim`") if d1_is_constant: return d2.rmin(d1) else: return d1.min(d2) def max_dim(d1: DimSize, d2: DimSize) -> DimSize: """Like max(d1, d2) but for both constant and symbolic dimensions.""" d1_is_constant = is_constant_dim(d1) if d1_is_constant and is_constant_dim(d2): return max(d1, d2) d1 = concrete_dim_or_error(d1, "argument `d1` of `core.max_dim`") d2 = concrete_dim_or_error(d2, "argument `d2` of `core.max_dim`") if d1_is_constant: return d2.rmax(d1) else: return d1.max(d2) def dimension_as_value(d: DimSize): """Turns a dimension size into a JAX array. This is the identity function for constant dimensions. Has the same abstract value as Python constants. """ if isinstance(d, (int, Tracer, np.int32, np.int64)): return d # For shape_poly._DimPolynomial m = getattr(d, "dimension_as_value", None) if m is not None: return m() return operator.index(d) def canonicalize_slice( s: slice, axis_size: DimSize ) -> tuple[DimSize, DimSize, DimSize]: """Computes the start index, step, and size of the slice `x[s]`. This is similar to `s.indices(axis_size)`, except that it returns `(start, step, size)`, and it works when the slice and/or the `axis_size` are symbolic. See https://numpy.org/doc/stable/user/basics.indexing.html#slicing-and-striding """ def convert_to_index(d: DimSize) -> DimSize: # Convert np.array and jax.Array to int, leave symbolic dimensions alone try: return operator.index(d) except: return d # Must resolve statically if step is {<0, ==0, >0} step = convert_to_index(s.step) if s.step is not None else 1 try: if step == 0: raise ValueError("slice step cannot be zero") step_gt_0 = (step > 0) except InconclusiveDimensionOperation as e: raise InconclusiveDimensionOperation( f"In slice with non-constant elements the step ({step}) must " + f"be resolved statically if it is > 0 or < 0.\nDetails: {e}") def clamp_index(i: DimSize, which: str): try: i_ge_0 = (i >= 0) except InconclusiveDimensionOperation as e: raise InconclusiveDimensionOperation( f"In slice with non-constant elements the {which} ({i}) must " + f"be resolved statically if it is >= 0.\nDetails: {e}") if i_ge_0: if step_gt_0: return min_dim(axis_size, i) else: return min_dim(axis_size - 1, i) else: if step_gt_0: return max_dim(0, axis_size + i) else: return max_dim(-1, axis_size + i) if s.start is None: start = 0 if step_gt_0 else axis_size - 1 else: start = clamp_index(convert_to_index(s.start), "start") if s.stop is None: stop = axis_size if step_gt_0 else -1 else: stop = clamp_index(convert_to_index(s.stop), "stop") gap = step if step_gt_0 else - step distance = (stop - start) if step_gt_0 else (start - stop) slice_size = max_dim(0, distance + gap - 1) // gap return start, step, slice_size class SomeTracer: __slots__ = () def __repr__(self): return "[dynamic]" def replace_tracer_for_error_message(obj): # TODO(mattjj): Many ideas for improving this. Crawl the stack and see if # there are user variables whose value is == to this object? Or search # parameters of functions being transformed, at least? Or at least assign # short unique ids to them? if isinstance(obj, Tracer): return SomeTracer() else: return obj def evaluate_shape(shape: Shape, dim_vars: Sequence[str], *dim_values: Array) -> Sequence[Array]: """Evaluates a shape possibly containing non-constants. Args: shape: the shape to evaluate. dim_vars: the dimension variables names that may appear in `shape`. dim_values: the dimension values corresponding to `dim_vars`. Returns: a tuple of JAX values corresponding to `shape`, of type `dim_value_dtype`. """ env = dict(zip(dim_vars, dim_values)) def eval_one_dim(d: DimSize): try: return operator.index(d) except: # Is a _DimExpr return d._evaluate(env) # pyrefly: ignore[missing-attribute] return tuple(eval_one_dim(d) for d in shape) def dim_value_dtype(): """The dtype to be used for dimension values.""" return dtypes.default_int_dtype() def dim_constant(ct: int): dtype = dim_value_dtype() assert dtype in (np.int32, np.int64) if dtype == np.int32: return np.int32(ct) elif dtype == np.int64: return np.int64(ct) def dim_value_aval() -> AbstractValue: return ShapedArray((), dim_value_dtype(), weak_type=True, sharding=None) # ------------------- Call ------------------- # eval_jaxpr_p is a call-like primitive parameterized by a jaxpr rather than a # Python callable: applying it evaluates the jaxpr, and staging it out is O(1), # emitting a single eqn that keeps its identity under retracing. Its # transformation rules live in partial_eval.py and lax/eval_jaxpr.py. eval_jaxpr_p = Primitive('eval_jaxpr') eval_jaxpr_p.multiple_results = True eval_jaxpr_p.def_impl(lambda *args, call_jaxpr, **_: jaxpr_as_fun(call_jaxpr)(*args)) eval_jaxpr_p.def_effectful_abstract_eval( lambda *_, call_jaxpr, **__: (call_jaxpr.out_avals, positional_effects(call_jaxpr))) # Aliases for the deleted final-style call primitives, for downstream code that # matches on these names when interpreting jaxprs. call_p = closed_call_p = eval_jaxpr_p # ------------------- Map ------------------- def mapped_aval(size: AxisSize, axis, aval: AbstractValue) -> AbstractValue: from jax._src.hijax import HiType # pyrefly: ignore[missing-import] if isinstance(aval, HiType): return aval.dec_rank(size, axis) # pyrefly: ignore[bad-argument-type] handler, _ = aval_mapping_handlers.get(type(aval), (None, None)) if handler is not None: return handler(size, axis, aval) else: raise TypeError(f"no mapping handler for {aval} of type {type(aval)}") def mapped_leading_aval(size, aval) -> AbstractValue: return mapped_aval(size, aval.leading_axis_spec(), aval) # TODO(yashkatariya): take axis data def unmapped_aval(size: AxisSize, axis: int | None, aval: AbstractValue, explicit_mesh_axis=None) -> AbstractValue: from jax._src.hijax import HiType # pyrefly: ignore[missing-import] if isinstance(aval, HiType): return aval.inc_rank(size, axis) # pyrefly: ignore[bad-argument-type] _, handler = aval_mapping_handlers.get(type(aval), (None, None)) if handler is not None: return handler(size, axis, explicit_mesh_axis, aval) else: raise TypeError(f"no unmapping handler for {aval} of type {type(aval)}") def unmapped_leading_aval(size, aval) -> AbstractValue: return unmapped_aval(size, aval.leading_axis_spec(), aval) def _map_shaped_array( size: int, axis: int | None, aval: ShapedArray) -> ShapedArray: assert axis is None or aval.shape[axis] == size if axis is None: return aval aval_s = aval.sharding sharding = aval_s.update( spec=aval_s.spec.update(partitions=tuple_delete(aval_s.spec.partitions, axis))) return ShapedArray(tuple_delete(aval.shape, axis), aval.dtype, weak_type=aval.weak_type, sharding=sharding, manual_axis_type=aval.mat, memory_space=aval.memory_space) def _unmap_shaped_array( size: int, axis: int | None, explicit_mesh_axis, aval: ShapedArray ) -> ShapedArray: if axis is None: return aval elif type(axis) is int: aval_s = aval.sharding sharding = aval_s.update(spec=aval_s.spec.update(partitions=tuple_insert( aval_s.spec.partitions, axis, explicit_mesh_axis))) return ShapedArray(tuple_insert(aval.shape, axis, size), aval.dtype, weak_type=aval.weak_type, sharding=sharding, manual_axis_type=aval.mat, memory_space=aval.memory_space) else: raise TypeError(axis) AvalMapHandlerPair = tuple[Callable, Callable] aval_mapping_handlers: dict[type, AvalMapHandlerPair] = { ShapedArray: (_map_shaped_array, _unmap_shaped_array), AbstractToken: (lambda _, __, a: a, lambda _, __, ____, a: a) } # When a mapped function is given no axis name, we generate a name object based # on the id of the function object. Collisions aren't important because this # name can't be used in collectives, as user code never gets a ref to this # object. We don't want to use the function object itself because that might # persist references to the function object. # TODO(mattjj): revisit this unique axis name strategy @total_ordering class _TempAxisName: def __init__(self, obj): self.id = id(obj) def __repr__(self): return f'' def __hash__(self): return hash(self.id) def __eq__(self, other): return type(other) is _TempAxisName and self.id == other.id def __lt__(self, other): return type(other) is _TempAxisName and self.id < other.id @dataclass(frozen=True, slots=True) class NamedAxisEffect(effects.Effect): """A side-effect introducing a new named axis into the current scope.""" name: AxisName effects.control_flow_allowed_effects.add_type(NamedAxisEffect) effects.custom_derivatives_allowed_effects.add_type(NamedAxisEffect) effects.lowerable_effects.add_type(NamedAxisEffect) effects.remat_allowed_effects.add_type(NamedAxisEffect) def filter_named_axis_effects( effects: Effects, names: Collection[AxisName] ) -> Effects: return {e for e in effects if not isinstance(e, NamedAxisEffect) or e.name not in names} def remove_named_axis_effects( jaxpr: Jaxpr, names: Collection[AxisName] ) -> Jaxpr: if not names or not jaxpr.effects: return jaxpr return jaxpr.replace(effects=filter_named_axis_effects(jaxpr.effects, names)) def replace_jaxpr_effects(jaxpr: Jaxpr, effects: Effects): return _replace_jaxpr_effects(jaxpr, frozenset(effects)) @weakref_lru_cache def _replace_jaxpr_effects(jaxpr: Jaxpr, effects: frozenset[Effect]): return jaxpr.replace(effects=set(effects)) # ------------------- Jaxpr checking ------------------- def typecheck(aval: AbstractValue, x) -> bool: return typecompat(aval, typeof(x)) def typecompat(aval_ref: AbstractValue, aval: AbstractValue) -> bool: """Determine whether `aval` conforms to `aval_ref`. Ignores weak_type.""" try: return typematch(aval_ref, aval) except TypeError: return False def typematch(t1: AbstractValue, t2: AbstractValue, no_dtype_check: bool = False) -> bool: """Determine whether `t1` and `t2` are equivalent. Ignores weak_type.""" t1 = t1.normalize() t2 = t2.normalize() from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] if t1 == t2: return True elif isinstance(t1, ShapedArray) and isinstance(t2, ShapedArray): if no_dtype_check: return cmp_shape_shd_mat_memsp(t1, t2) return t1.dtype == t2.dtype and cmp_shape_shd_mat_memsp(t1, t2) elif isinstance(t1, AbstractRef) and isinstance(t2, AbstractRef): # We want to use the regular typecheck for ShapedArray here. return (typematch(t1.inner_aval, t2.inner_aval, no_dtype_check) and (t1.memory_space is None or t2.memory_space is None or t1.memory_space == t2.memory_space)) else: return False def cmp_shape_shd_mat_memsp(t1, t2): # TODO(yashkatariya): Expand this to Manual and Auto mode. # See https://github.com/jax-ml/jax/issues/26474 t1_mesh, t2_mesh = t1.sharding.mesh, t2.sharding.mesh if not t1_mesh.empty and not t2_mesh.empty: if t1_mesh._any_axis_explicit or t2_mesh._any_axis_explicit: shd_eq = t1.sharding == t2.sharding else: shd_eq = True else: shd_eq = True return (shd_eq and definitely_equal_shape(t1.shape, t2.shape) and t1.mat == t2.mat and t1.memory_space == t2.memory_space) def aval_mismatch_extra(a1: AbstractValue, a2: AbstractValue) -> str: assert not typematch(a1, a2) if isinstance(a1, ShapedArray) and isinstance(a2, ShapedArray): mismatches = [] if a1.dtype != a2.dtype: mismatches.append('the dtypes do not match') if a1.shape != a2.shape: mismatches.append('the shapes do not match') if a1.mat != a2.mat: mismatches.append('the manual axis types do not match') # TODO(yashkatariya,mattjj): add check for sharding-in-types mismatch if len(mismatches) == 0: return '' elif len(mismatches) == 1: return ', so ' + mismatches[0] else: return ', so ' + ', '.join(mismatches[:-1]) + ', and ' + mismatches[-1] return '' @set_module("jax.errors") class JaxprTypeError(TypeError): pass custom_typechecks: dict[Primitive, Callable] = {} def _check_closed_call(_, *in_atoms, call_jaxpr, **__): in_avals = [x.aval for x in in_atoms] if not all(map(typecompat, call_jaxpr.in_avals, in_avals)): raise JaxprTypeError("Closed call in_avals mismatch") return call_jaxpr.out_avals, positional_effects(call_jaxpr) custom_typechecks[eval_jaxpr_p] = _check_closed_call def check_jaxpr(jaxpr: Jaxpr): """Checks well-formedness of a jaxpr. Specifically, check that: - variables that are read are bound beforehand - variables are typed equally throughout a jaxpr - variable type annotations are compatible with their binding expression Raises `JaxprTypeError` if `jaxpr` is determined invalid. Returns `None` otherwise. """ @functools.cache def ctx_factory(): ctx = JaxprPpContext(_dropvars(jaxpr)) pp_settings = JaxprPpSettings() try: pp_jaxpr(jaxpr, ctx, pp_settings) # side-effect on ctx, build variable names except: pass return ctx, pp_settings try: _check_jaxpr(ctx_factory, jaxpr) except JaxprTypeError as e: ctx, pp_settings = ctx_factory() if len(e.args) == 2: msg, eqnidx = e.args jaxpr_str = str(pp_jaxpr_eqn_range(jaxpr, eqnidx - 10, eqnidx + 10, ctx, pp_settings)) else: msg, = e.args jaxpr_str = str(pp_jaxpr_eqn_range(jaxpr, 0, 20, ctx, pp_settings)) msg = "\n\n".join([msg, "while checking jaxpr:", jaxpr_str]) raise JaxprTypeError(msg) from None # Run key reuse checker after validating jaxpr: if config.debug_key_reuse.value: # Import here to avoid circular imports from jax.experimental.key_reuse._core import check_key_reuse_jaxpr # pyrefly: ignore[missing-import] check_key_reuse_jaxpr(jaxpr) @partial(weakref_lru_cache, trace_context_in_key=False) def _dropvars(jaxpr: Jaxpr) -> dict[Var, Literal_['_']]: varnames: dict[Var, Literal_['_']] = {} used: set[Var] = {atom for atom in jaxpr.outvars if isinstance(atom, Var)} for eqn in jaxpr.eqns[::-1]: for v in eqn.outvars: if not v in used: varnames[v] = '_' used.update(atom for atom in eqn.invars if isinstance(atom, Var)) return varnames def _check_jaxpr( ctx_factory: Callable[[], tuple[JaxprPpContext, JaxprPpSettings]], jaxpr: Jaxpr ) -> None: env: dict[Var, Atom] = {} def read(x: Atom) -> Atom: # Check the type annotation is itself well-typed. check_type(ctx_factory, env, x.aval) if isinstance(x, Var): # Check the variable is in-scope and consistently typed. if x not in env: ctx, _ = ctx_factory() raise JaxprTypeError(f"Variable '{x.pretty_print(ctx)}' not defined") return env[x] elif isinstance(x, Literal): # Check that the literal matches its type annotation. if not typecheck(x.aval, x.val): ctx, _ = ctx_factory() raise JaxprTypeError( f"Literal value {x.val} does not match its type annotation " f"{pp_aval(x.aval, ctx)}") return x else: assert False, "syntactically invalid jaxpr" def write(v: Var, aval: AbstractValue) -> None: assert isinstance(v, Var), "syntactically invalid jaxpr" # Check the type annotation of the binder is itself well-typed. check_type(ctx_factory, env, v.aval) # Check that the variable is not already bound. if v in env: ctx, _ = ctx_factory() raise JaxprTypeError(f"Variable '{v.pretty_print(ctx)}' already bound") # Check that the computed type is consistent with the binder annotation. if not typematch(v.aval, aval): ctx, _ = ctx_factory() raise JaxprTypeError( f"Value for variable '{v.pretty_print(ctx)}' inconsistently typed " f"as {pp_aval(aval, ctx)} for let-binder of type {pp_aval(v.aval, ctx)}") # If the variable is not a DropVar, add it to the environment. if not isinstance(v, DropVar): env[v] = v # # Don't return refs if config.mutable_array_checks.value: from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] for v in jaxpr.outvars: if isinstance(v.aval, AbstractRef): raise JaxprTypeError("returned a ref!") # Check type annotations on lambda binders. for v in it.chain(jaxpr.constvars, jaxpr.invars): check_type(ctx_factory, env, v.aval) write(v, v.aval) # Check each eqn. input_vars = set(it.chain(jaxpr.constvars, jaxpr.invars)) mut_arrays = set() for eqn_idx, eqn in enumerate(jaxpr.eqns): prim = eqn.primitive try: in_atoms = map(read, eqn.invars) in_avals = [x.aval for x in in_atoms] # use in_atoms for dyn shapes # Compute the type of the primitive application. with eqn.ctx.manager: if prim in custom_typechecks: out_type, eqn_effects = custom_typechecks[prim]( ctx_factory, *in_atoms, **eqn.params) else: out_type, eqn_effects = check_eqn(prim, in_avals, eqn.params) # Check the computed effect type matches the eqn's annotation, and is # included in the jaxpr's annotation. if prim.ref_primitive: if prim.ref_allocating: outvar, = eqn.outvars mut_arrays.add(outvar) eqn_effects = resolve_input_effects(eqn_effects, eqn.invars) if eqn.effects != eqn_effects: raise JaxprTypeError("Inferred effects do not match equation effects. " f"Equation effects: {eqn.effects}. " f"Inferred effects: {eqn_effects}") for eff in eqn.effects: if isinstance(eff, effects.JaxprInputEffect): if eff.input in mut_arrays: continue if eff.input not in input_vars: raise JaxprTypeError( "Invalid `JaxprInputEffect`: must correspond to a jaxpr invar") if eff not in jaxpr.effects: raise JaxprTypeError( "Invalid `JaxprInputEffect`: must be present in jaxpr. " f"{eff} is not in {jaxpr.effects}.") elif isinstance(eff, NamedAxisEffect): # It is valid for a primitive to discharge the named axis effect. continue elif eff not in jaxpr.effects: raise JaxprTypeError("Equation effect not present in jaxpr effects. " f"Equation effect: {eff}. " f"Jaxpr effects: {jaxpr.effects}") # Check out_type matches the let-binders' annotation (after substitution). foreach(write, eqn.outvars, out_type) except JaxprTypeError as e: ctx, settings = ctx_factory() msg, = e.args src = source_info_util.summarize(eqn.source_info) msg = "\n\n".join([msg, "in equation:", str(pp.nest(2, pp_eqn(eqn, ctx, settings))), f"from source: {src}"]) raise JaxprTypeError(msg, eqn_idx) from None # Check there are no output refs # TODO(mattjj): improve this error message if config.mutable_array_checks.value: from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] for v in jaxpr.outvars: if isinstance(v.aval, AbstractRef): raise TypeError("returned ref") # TODO(mattjj): include output type annotation on jaxpr and check it here foreach(read, jaxpr.outvars) def check_type( ctx_factory: Callable[[], tuple[JaxprPpContext, JaxprPpSettings]], env: dict[Var, Atom], ty: AbstractValue, ) -> None: return # Except in above case(s), all syntactic forms are valid def check_eqn(prim, in_avals, params): for jaxpr in jaxprs_in_params(params): check_jaxpr(jaxpr) out_avals, effects = prim.abstract_eval(*in_avals, **params) if not prim.multiple_results: out_avals = [out_avals] return out_avals, effects def _check_call(ctx_factory, prim, in_atoms, params): if "call_jaxpr" not in params: raise JaxprTypeError( f"Call primitive {prim} missing 'call_jaxpr' parameter") call_jaxpr = params["call_jaxpr"] if len(in_atoms) != len(call_jaxpr.invars): raise JaxprTypeError(f"Call primitive {prim} with {len(in_atoms)} " f"operands cannot call jaxpr with " f"{len(call_jaxpr.invars)} inputs") # Check `call_jaxpr` can be applied to in_atoms. env: dict[Var, Atom] = {} for v, x in zip(call_jaxpr.invars, in_atoms): if not typecompat(v.aval, x.aval): # TODO(mattjj): vars in error message are confusing b/c of Var.__repr__ raise JaxprTypeError(f"Call primitive {prim} passes operand {x} of type " f"{x.aval} to jaxpr expecting type " f"{v.aval}") env[v] = x.val if type(x) is Literal else x check_jaxpr(call_jaxpr) out_avals = [x.aval for x in call_jaxpr.outvars] return out_avals, positional_effects(call_jaxpr) def eqn_effects(jaxpr, invars) -> Effects: return resolve_input_effects(positional_effects(jaxpr), invars) def subst_input_effects(effs, env) -> Effects: return {e.replace(env.get(e.input, e.input)) if isinstance(e, effects.JaxprInputEffect) else e for e in effs} def positional_effects(jaxpr) -> Effects: if not any(isinstance(e, effects.JaxprInputEffect) for e in jaxpr.effects): return jaxpr.effects idx = {v: i for i, v in enumerate(jaxpr.invars)} out_effs = set() for eff in jaxpr.effects: if isinstance(eff, effects.JaxprInputEffect): i = idx.get(eff.input) if i is None: continue eff = eff.replace(i) out_effs.add(eff) return out_effs # ------------------- ShapeDtypeStruct ------------------- def _check_sharding(sharding, shape): if sharding is None: return if isinstance(sharding, P): sharding._check_compatible_wrt_shape(shape) else: sharding.check_compatible_aval(shape) @set_module("jax") class ShapeDtypeStruct: """A container for the shape, dtype, and other static attributes of an array. ``ShapeDtypeStruct`` is often used in conjunction with :func:`jax.eval_shape`. Args: shape: a sequence of integers representing an array shape dtype: a dtype-like object sharding: (optional) a :class:`jax.Sharding` object """ __slots__ = ["shape", "dtype", "_sharding", "_dll", "weak_type", "manual_axis_type", "is_ref", "_memory_space"] shape: Any dtype: Any _sharding: Any _dll: Any weak_type: Any manual_axis_type: Any is_ref: Any _memory_space: Any def __init__(self, shape, dtype, *, sharding=None, weak_type=False, manual_axis_type=None, is_ref=False, _memory_space=None): shape = tuple(shape) if any(s is None for s in shape): raise ValueError('`shape` passed to `ShapeDtypeStruct` cannot have ' f'None in it. Got {shape=}') object.__setattr__(self, 'shape', shape) if dtype is None: raise ValueError("ShapeDtypeStruct: dtype must be specified.") dtype = dtype if dtypes.issubdtype(dtype, dtypes.extended) else np.dtype(dtype) object.__setattr__(self, 'dtype', dtype) if sharding is not None and not isinstance(sharding, (Sharding, Format, P)): raise ValueError( "sharding should be an instance of `jax.sharding.Sharding`, " "`jax.sharding.PartitionSpec` or" f" `jax.experimental.layout.Format`. Got {sharding} of type" f" {type(sharding)}.") if (isinstance(sharding, Format) and isinstance(sharding.layout, AutoLayoutSingleton)): raise TypeError( "`Layout.AUTO` cannot be used in place of a device-local" f" layout in a `ShapeDtypeStruct`. Got {sharding}") object.__setattr__( self, '_dll', sharding.layout if isinstance(sharding, Format) else None) object.__setattr__( self, '_sharding', sharding.sharding if isinstance(sharding, Format) else sharding) _check_sharding(self._sharding, self.shape) object.__setattr__(self, 'weak_type', weak_type) if (manual_axis_type is not None and not isinstance(manual_axis_type, ManualAxisType)): raise TypeError( "`manual_axis_type` argument passed to ShapeDtypeStruct should be of" " type `jax.sharding.ManualAxisType`. Got type" f" {type(manual_axis_type)}") object.__setattr__(self, 'manual_axis_type', manual_axis_type) object.__setattr__(self, 'is_ref', is_ref) object.__setattr__(self, '_memory_space', _memory_space) def __setattr__(self, name, value): if hasattr(self, name): if getattr(self, name) == value: # This can happen if two threads race, for example if two threads # are trying to hash the same SDS instance. return raise RuntimeError( f"Cannot reassign attributes ({name}) of immutable ShapeDtypeStruct" " objects") super().__setattr__(name, value) size = property(lambda self: math.prod(self.shape)) ndim = property(lambda self: len(self.shape)) @classmethod def like(cls, x) -> ShapeDtypeStruct: from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] from jax._src.array import ArrayImpl # pyrefly: ignore[missing-import] from jax._src.hijax import HiType # pyrefly: ignore[missing-import] aval = shaped_abstractify(x) if isinstance(aval, AbstractRef): raise NotImplementedError( "Ref/AbstractRef support is not implemented. Please file an issue at" " https://github.com/jax-ml/jax/issues") if isinstance(aval, HiType): raise NotImplementedError( "Passing a HiVal to `ShapeDtypeStruct.like` is not implemented.") aval_s = None if aval.sharding.mesh.empty else aval.sharding if getattr(x, '_committed', True): sharding = (x.format if isinstance(x, ArrayImpl) else getattr(x, 'sharding', aval_s)) else: sharding = None mat = None if aval.mat.empty else aval.mat return cls(aval.shape, aval.dtype, sharding=sharding, weak_type=aval.weak_type, manual_axis_type=mat, is_ref=False, _memory_space=aval.memory_space) @property def sharding(self): if isinstance(self._sharding, P): # TODO(yashkatariya): Maybe use `get_abstract_mesh()` here but switch # on `core.trace_state_clean()`? cur_mesh = mesh_lib.get_concrete_mesh() if cur_mesh.empty: raise TypeError( "When specifying PartitionSpec to `ShapeDtypeStruct`, the context" " mesh cannot be empty. Please use `jax.set_mesh` to set" " the mesh context.") return NamedSharding(cur_mesh, self._sharding) else: return self._sharding @property def format(self): return Format(self._dll, self.sharding) def __len__(self): try: return self.shape[0] except IndexError as e: raise TypeError("len() of unsized object") from e # same as numpy error def __repr__(self): sh = f", sharding={self.sharding}" if self.sharding is not None else "" l = f", format={self._dll}" if self._dll is not None else "" wt = f", weak_type={self.weak_type}" if self.weak_type else "" mat = (f", manual_axis_type={self.manual_axis_type}" if self.manual_axis_type else "") is_ref = f", is_ref={self.is_ref}" if self.is_ref else "" return (f"{type(self).__name__}(shape={self.shape}, " f"dtype={self.dtype.name}{sh}{l}{wt}{mat}{is_ref})") __str__ = __repr__ def __eq__(self, other): if not isinstance(other, ShapeDtypeStruct): return False else: return ((self.shape, self.dtype, self.sharding, self._dll, self.weak_type, self.manual_axis_type, self.is_ref) == (other.shape, other.dtype, other.sharding, other._dll, other.weak_type, other.manual_axis_type, other.is_ref)) def __hash__(self): return hash((self.shape, self.dtype, self.sharding, self._dll, self.weak_type, self.manual_axis_type, self.is_ref)) def update(self, **kwargs): if 'sharding' in kwargs: s = kwargs['sharding'] if self._dll is not None and isinstance(s, Sharding): raise ValueError( f"You are updating ShapeDtypeStruct with a {type(s)} when the" f" original ShapeDtypeStruct had a concrete layout {self.format}." " This might lead to bugs. If you want to do this, create a new" " ShapeDtypeStruct via the constructor.") sharding = s else: sharding = self.format return ShapeDtypeStruct( shape=kwargs.pop('shape', self.shape), dtype=kwargs.pop('dtype', self.dtype), sharding=sharding, weak_type=kwargs.pop('weak_type', self.weak_type), manual_axis_type=kwargs.pop('manual_axis_type', self.manual_axis_type), is_ref=kwargs.pop('is_ref', self.is_ref), _memory_space=kwargs.pop('_memory_space', self._memory_space)) def _sds_aval_mapping(x): dtype = dtypes.check_and_canonicalize_user_dtype( x.dtype, "ShapeDtypeStruct", allow_non_jax_dtypes=True) memory_space = getattr(x, "_memory_space", None) or MemorySpace.Device aval = ShapedArray(x.shape, dtype, weak_type=x.weak_type, memory_space=memory_space) aval = update_aval_with_sharding(aval, x.sharding, mat=x.manual_axis_type) if x.is_ref: from jax._src.state.types import AbstractRef # pyrefly: ignore[missing-import] return AbstractRef(aval) return aval def _canonicalize_sds(sds): dt = dtypes.check_and_canonicalize_user_dtype( sds.dtype, "ShapeDtypeStruct", allow_non_jax_dtypes=True) if dt == sds.dtype: return sds return sds.update(dtype=dt) pytype_aval_mappings[ShapeDtypeStruct] = _sds_aval_mapping dtypes.register_canonicalize_value_handler(ShapeDtypeStruct, _canonicalize_sds) # ------------------- Jaxpr printed representation ------------------- def pp_toplevel_jaxpr(jaxpr_to_print: Jaxpr, *, source_info: bool = False, print_shapes: bool = True, custom_pp_eqn_rules : bool = True, name_stack: bool = False, print_effects: bool = False) -> pp.Doc: context = JaxprPpContext(_dropvars(jaxpr_to_print)) settings = JaxprPpSettings( source_info=source_info, print_shapes=print_shapes, custom_pp_eqn_rules=custom_pp_eqn_rules, name_stack=name_stack, print_effects=print_effects) # Compute how many times each jaxpr is used. names = defaultdict[Jaxpr, str](lambda: "jaxpr") jaxpr_counts = Counter[Jaxpr]() s = deque([jaxpr_to_print]) while s: jaxpr = s.popleft() jaxpr_counts[jaxpr] += 1 if jaxpr is not jaxpr_to_print and len(jaxpr.eqns) > 10: jaxpr_counts[jaxpr] += 1 for eqn in jaxpr.eqns: for name, subjaxpr in _pp_subjaxprs(eqn, settings): s.append(subjaxpr) if name is not None: names.setdefault(subjaxpr, name.strip("<>")) # -> lambda # Pull jaxprs occurring more than once to the top-level, making sure # that their names are unique. name_counts = Counter[str]() shared = [] for jaxpr, c in jaxpr_counts.items(): if c == 1: continue name = names[jaxpr] if (count := name_counts[name]) > 0: name_counts[name] += 1 name += str(count) name_counts[name] += 1 else: name_counts[name] += 1 context.shared_jaxpr_names.add(name) context.shared_jaxprs[jaxpr] = name shared.append((name, jaxpr)) docs = [] for name, jaxpr in shared: docs.append(pp_shared_jaxpr(name, jaxpr, context, settings)) docs.append(pp_jaxpr(jaxpr_to_print, context, settings)) return pp.concat(docs) # Rules giving the jaxprs an eqn prints, each paired with a name (or None) to use # if it's hoisted as a shared jaxpr, for primitives that print jaxprs other than # their Jaxpr-valued params (e.g. hijax primitives, which nest them in `_prim`). pp_subjaxprs_rules: dict[ Primitive, Callable[[JaxprEqn], Iterable[tuple[str | None, Jaxpr]]]] = {} def _pp_subjaxprs(eqn: JaxprEqn, settings: JaxprPpSettings ) -> Iterable[tuple[str | None, Jaxpr]]: if settings.custom_pp_eqn_rules and eqn.primitive in pp_subjaxprs_rules: return pp_subjaxprs_rules[eqn.primitive](eqn) # TODO(slebedev): Come up with a more elaborate heuristic for name=. name = eqn.params.get("name") return [(name, j) for j in jaxprs_in_params(eqn.params)] class JaxprPpSettings(NamedTuple): print_shapes: bool = True source_info: bool = False name_stack: bool = False custom_pp_eqn_rules: bool = True print_effects: bool = False def _encode_digits_alphabetic(n: int) -> str: if n == -1: return '*' s = '' while len(s) == 0 or n: n, i = n // 26, n % 26 s = chr(97 + i % 26) + s return s # A JaxprPpContext allows globally unique variable names within nested Jaxprs. class JaxprPpContext: var_names: defaultdict[Var, str] # Shared jaxprs are those that are used multiple times and are printed first. shared_jaxprs: MutableMapping[Jaxpr, str] # maps shared jaxpr to its name shared_jaxpr_names: MutableSet[str] def __init__(self, var_names: dict | None = None) -> None: self.shared_jaxprs = {} self.shared_jaxpr_names = set() fresh_names: Iterator[str] = ( name for i in it.count() if (name := _encode_digits_alphabetic(i)) not in self.shared_jaxpr_names) self.var_names = defaultdict(fresh_names.__next__, var_names or {}) def suggest_same_var_names(self, for_vars: Sequence[Atom], like_vars: Sequence[Atom]) -> None: """Suggests the names for `for_vars` to match those of `like_vars`. `for_vars` are distinct Vars, and are aliased with `like_vars`. """ used_like_vars: set[Var] = set() if len(for_vars) != len(like_vars): # The mismatch can happen if a primitive containing a subjaxpr is invoked # with the wrong number of arguments, e.g., when printing an invalid Jaxpr. return for for_v, like_v in zip(for_vars, like_vars): if (isinstance(like_v, Var) and like_v not in used_like_vars and isinstance(for_v, Var) and for_v not in self.var_names): used_like_vars.add(like_v) self.var_names[for_v] = like_v.pretty_print(self) def pp_var(v: Var | Literal, context: JaxprPpContext, *, print_literal_dtype: bool = True, is_binder: bool = False) -> pp.Doc: name = v.pretty_print(context, print_dtype=print_literal_dtype) if (isinstance(v, Var) and not isinstance(v, DropVar)): if is_binder: return pp.text(name, anchor=f"v_{name}") else: return pp.text(name, href=f"#v_{name}") return pp.text(name) def pp_aval(a: AbstractValue, context: JaxprPpContext) -> str: return a.str_short(short_dtypes=True) def pp_vars(vs: Sequence[Atom], context: JaxprPpContext, *, separator="", print_shapes: bool = False, is_binder: bool = False) -> pp.Doc: if print_shapes: return pp.nest(2, pp.group( pp.join(pp.text(separator) + pp.group(pp.brk()), [ pp_var(v, context, is_binder=is_binder) + pp.type_annotation(pp.text(":" + pp_aval(v.aval, context))) for v in vs ]) )) else: return pp.nest(2, pp.group( pp.join(pp.text(separator) + pp.group(pp.brk()), [pp_var(v, context, is_binder=is_binder) for v in vs]) )) def pp_kv_pair(k:str, v: Any, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: if type(v) is tuple and all(isinstance(j, Jaxpr) for j in v): pp_v = pp_jaxprs(v, context, settings) elif isinstance(v, Jaxpr): pp_v = pp_jaxpr(v, context, settings) elif isinstance(v, frozenset): pp_v = pp.text(f"frozenset({{{', '.join(repr(e) for e in sorted(v))}}})") elif isinstance(v, pp.Doc): pp_v = v else: s = str(v) s = re.sub( r' at 0x([0-9a-fA-F]+)', lambda m: ' at 0x' + 'X' * len(m.group(1)), s) pp_v = pp.text(s) return pp.text(f'{k}=') + pp_v def pp_kv_pairs(kv_pairs, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: if not kv_pairs: return pp.nil() return pp.group(pp.concat([ pp.nest(2, pp.concat([ pp.text("["), pp.brk(""), pp.join(pp.brk(), [pp_kv_pair(k, v, context, settings) for k, v in kv_pairs]) ])), pp.brk(""), pp.text("]") ])) def pp_eqn(eqn: JaxprEqn, context: JaxprPpContext, settings: JaxprPpSettings ) -> pp.Doc: rule = (_pp_eqn if not settings.custom_pp_eqn_rules else pp_eqn_rules.get(eqn.primitive, _pp_eqn)) doc = rule(eqn, context, settings) return (doc if eqn.source_info.traceback is None else pp.source_map(doc, eqn.source_info.traceback)) def _pp_eqn(eqn: JaxprEqn, context: JaxprPpContext, settings: JaxprPpSettings, params: Sequence[str] | None = None, name: str | None = None ) -> pp.Doc: annotation = (source_info_util.summarize(eqn.source_info) if settings.source_info else None) if params is None: params = sorted(eqn.params) name_stack_annotation = f'[{eqn.source_info.name_stack}]' if settings.name_stack else None lhs = pp_vars(eqn.outvars, context, print_shapes=settings.print_shapes, is_binder=True) rhs = [pp.text(name or eqn.primitive.name, annotation=name_stack_annotation), pp_kv_pairs([(p, eqn.params[p]) for p in params], context, settings), pp.text(" ") + pp_vars(eqn.invars, context)] if eqn.outvars: return pp.concat([lhs, pp.text(" = ", annotation=annotation), *rhs]) else: return pp.concat(rhs) CustomPpEqnRule = Callable[[JaxprEqn, JaxprPpContext, JaxprPpSettings], pp.Doc] pp_eqn_rules: dict[Primitive, CustomPpEqnRule] = {} def pp_eqns(eqns: Sequence[JaxprEqn], context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: return pp.join( pp.brk("; "), [pp_eqn(e, context, settings) for e in eqns]) def pp_jaxpr_skeleton(jaxpr: Jaxpr, eqns_fn, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: constvars = pp_vars(jaxpr.constvars, context, print_shapes=settings.print_shapes, is_binder=True) invars = pp_vars(jaxpr.invars, context, print_shapes=settings.print_shapes, is_binder=True) eqns = eqns_fn() outvars = pp.concat([ pp.text("("), pp_vars(jaxpr.outvars, context, separator=","), pp.text(")" if len(jaxpr.outvars) != 1 else ",)")]) if settings.print_effects: # TODO(sharadmv): render an entire signature here eff_text = [pp.text(" : { ")] for i, eff in enumerate(jaxpr.effects): if i > 0: eff_text.append(pp.text(", ")) eff_text.append(pp_effect(eff, context)) eff_text.append(pp.text(" }")) else: eff_text = [] return pp.group(pp.nest(2, pp.concat([ pp.text("{ "), pp.keyword(pp.text("lambda ")), constvars, pp.text("; "), invars, pp.text(". "), pp.keyword(pp.text("let")), pp.nest(2, pp.brk() + eqns), pp.brk(), pp.keyword(pp.text("in ")), outvars, pp.concat(eff_text) ])) + pp.text(" }")) def pp_shared_jaxpr( name: str, jaxpr: Jaxpr, context: JaxprPpContext, settings: JaxprPpSettings, ) -> pp.Doc: eqns_fn = lambda: pp_eqns(jaxpr.eqns, context, settings) pp_skeleton = pp_jaxpr_skeleton(jaxpr, eqns_fn, context, settings) return pp.concat([ pp.text("let "), pp.text(name, anchor=f"g_{name}"), pp.text(" = "), pp_skeleton, pp.text(" in"), pp.brk(), ]) def pp_jaxpr( jaxpr: Jaxpr, context: JaxprPpContext, settings: JaxprPpSettings, ) -> pp.Doc: if name := context.shared_jaxprs.get(jaxpr): return pp.text(name, href=f"#g_{name}") eqns_fn = lambda: pp_eqns(jaxpr.eqns, context, settings) return pp_jaxpr_skeleton(jaxpr, eqns_fn, context, settings) def pp_jaxprs( jaxprs: Sequence[Jaxpr], context: JaxprPpContext, settings: JaxprPpSettings ) -> pp.Doc: return pp.group(pp.concat([pp.nest(2, pp.concat([ pp.text('('), pp.brk(""), pp.join(pp.brk(), map(lambda x: pp_jaxpr(x, context, settings), jaxprs))] )), pp.brk(""), pp.text(')')]) ) def pp_jaxpr_eqn_range(jaxpr: Jaxpr, lo: int, hi: int, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: lo = max(lo, 0) hi = max(lo, min(hi, len(jaxpr.eqns))) eqns = jaxpr.eqns[lo:hi] def eqns_fn(): pps = [] if len(eqns) == 0 and len(jaxpr.eqns) != 0: pps.append(pp.text('...')) else: if lo != 0: pps.append(pp.text('...')) pps.extend(map((lambda e: pp_eqn(e, context, settings)), eqns)) if hi != len(jaxpr.eqns): pps.append(pp.text('...')) return pp.join(pp.brk("; "), pps) return pp_jaxpr_skeleton(jaxpr, eqns_fn, context, settings) def pp_effect(effect: Effect, context: JaxprPpContext) -> pp.Doc: if hasattr(effect, "_pretty_print"): return effect._pretty_print(context) return pp.text(str(effect)) # ------------------- Jaxpr util ------------------- def last_used(jaxpr: Jaxpr) -> dict[Var, JaxprEqn | None]: """Returns a mapping from every var in jaxpr to what equation uses it last.""" last_used: dict[Var, JaxprEqn | None] = { v: None for v in jaxpr.outvars if not isinstance(v, Literal)} for eqn in reversed(jaxpr.eqns): for v in eqn.invars: if not isinstance(v, Literal) and v not in last_used: last_used[v] = eqn return last_used def clean_up_dead_vars(eqn: JaxprEqn, env: dict[Var, Any], last_used: dict[Var, JaxprEqn | None]): """Remove all eqn.invars from env if eqn is the last time they were used.""" for v in {v for v in eqn.invars if not isinstance(v, Literal)}: if last_used[v] is eqn: # Delete ref to variable when it is no longer needed by next equations. del env[v] # Used in shard_map and reshard for converting avals shard_aval_handlers = {} unshard_aval_handlers = {} reshard_aval_handlers = {} def shard_aval(mesh, manual_axes, check_vma, spec, aval: AbstractValue ) -> AbstractValue: from jax._src.hijax import HiType # pyrefly: ignore[missing-import] if isinstance(aval, HiType): return aval.shard(mesh, manual_axes, check_vma, spec) if (handler := shard_aval_handlers.get(type(aval))): return handler(mesh, manual_axes, check_vma, spec, aval) raise NotImplementedError(f"Unsupported aval type: {type(aval)}") def unshard_aval(mesh, check_vma, spec, aval: AbstractValue ) -> AbstractValue: from jax._src.hijax import HiType # pyrefly: ignore[missing-import] if isinstance(aval, HiType): return aval.unshard(mesh, check_vma, spec) if (handler := unshard_aval_handlers.get(type(aval))): return handler(mesh, check_vma, spec, aval) raise NotImplementedError(f"Unsupported aval type: {type(aval)}") def reshard_aval(aval: AbstractValue, dst_sharding) -> AbstractValue: from jax._src.hijax import HiType # pyrefly: ignore[missing-import] if isinstance(aval, HiType): return aval.reshard(dst_sharding) if (handler := reshard_aval_handlers.get(type(aval))): return handler(aval, dst_sharding) raise NotImplementedError(f"Unsupported aval type: {type(aval)}") # ----------------- external APIs for querying tracing context ----------------- # TODO(dougalm, jakevdp): expose these via jax.extend # Comparable object for checking whether JAX's trace state has changed. class OpaqueTraceState: def __init__(self, trace_ref): self._trace_ref = trace_ref def __eq__(self, other): if isinstance(other, OpaqueTraceState): return self._trace_ref == other._trace_ref else: return False def get_opaque_trace_state(convention=None): del convention assert trace_ctx.trace is not None return OpaqueTraceState(trace_ctx.trace._weakref) def nonempty_axis_env() -> bool: return bool(trace_ctx.axis_env.axis_sizes) def unsafe_am_i_under_a_jit() -> bool: return 'DynamicJaxprTrace' in str(unsafe_get_trace_stack(trace_ctx.trace)) def unsafe_am_i_under_a_vmap() -> bool: return 'BatchTrace' in str(unsafe_get_trace_stack(trace_ctx.trace)) # TODO(douglam): deprecate/delete def find_top_trace(_): return unsafe_get_current_trace() def unsafe_get_current_trace(): return trace_ctx.trace def unsafe_get_trace_stack(trace): if hasattr(trace, "parent_trace"): return unsafe_get_trace_stack(trace.parent_trace) + [trace] else: return [trace] def unsafe_get_axis_names() -> list[Any]: return list(trace_ctx.axis_env.axis_sizes) # TODO(douglam): deprecate/delete def axis_frame(axis_name): return trace_ctx.axis_env.axis_size(axis_name)