Stackless [WIP]

.github/workflows/ci-build.yaml

@@ -210,4 +210,3 @@ jobs:
  echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
  echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
  pytest -n auto --tb=short --maxfail=20 jax/experimental/jax2tf/tests/jax2tf_test.py
-

jax/_src/ad_checkpoint.py

@@ -496,7 +496,8 @@ def print_saved_residuals(f, *args, **kwargs):
 @remat_p.def_impl
 def remat_impl(*args, jaxpr, prevent_cse, differentiated, policy):
  del prevent_cse, differentiated, policy # Unused.
- return core.eval_jaxpr(jaxpr, (), *args)
+ with core.concrete_eval():
+ return core.eval_jaxpr(jaxpr, (), *args)
 
 @remat_p.def_effectful_abstract_eval
 def remat_abstract_eval(*args, jaxpr, prevent_cse, differentiated, policy):
@@ -701,20 +702,17 @@ def transposed(*args_flat):
  transposed_jaxpr = core.ClosedJaxpr(transposed_jaxpr_, consts)
  return transposed_jaxpr, cell.in_cts_zero # pytype: disable=attribute-error
 
-def remat_vmap(spmd_axis_name, axis_size, axis_name, main_type, args, dims, *,
- jaxpr, **params):
+def remat_vmap(axis_data, args, dims, *, jaxpr, **params):
  assert not jaxpr.constvars
  jaxpr_batched_, out_batched = batching.batch_jaxpr_axes(
- pe.close_jaxpr(jaxpr), axis_size, dims,
- [batching.zero_if_mapped] * len(jaxpr.outvars),
- axis_name=axis_name, spmd_axis_name=spmd_axis_name, main_type=main_type)
+ pe.close_jaxpr(jaxpr), axis_data, dims,
+ [batching.zero_if_mapped] * len(jaxpr.outvars))
  jaxpr_batched, consts = jaxpr_batched_.jaxpr, jaxpr_batched_.consts
  if consts:
  jaxpr_batched = pe.convert_constvars_jaxpr(jaxpr_batched)
  out_dims = [0 if b else None for b in out_batched]
  return remat_p.bind(*consts, *args, jaxpr=jaxpr_batched, **params), out_dims
-batching.axis_primitive_batchers[remat_p] = partial(remat_vmap, None)
-batching.spmd_axis_primitive_batchers[remat_p] = remat_vmap
+batching.fancy_primitive_batchers[remat_p] = remat_vmap
 
 # TODO(mattjj,sharadmv): de-duplicate with pe.dce_jaxpr_call_rule
 def remat_dce(used_outputs: list[bool], eqn: core.JaxprEqn

jax/_src/ad_util.py

@@ -38,7 +38,8 @@ def add_jaxvals(x: ArrayLike, y: ArrayLike) -> Array:
 
 @add_jaxvals_p.def_impl
 def add_impl(x, y):
- return raw_jaxval_adders[type(x)](x, y)
+ with core.set_current_trace(core.EvalTrace()):
+ return raw_jaxval_adders[type(x)](x, y )
 raw_jaxval_adders = {} # type: ignore
 
 @add_jaxvals_p.def_abstract_eval

jax/_src/api.py

@@ -34,7 +34,7 @@
 import weakref
 
 import numpy as np
-from contextlib import contextmanager, ExitStack
+from contextlib import contextmanager
 
 from jax._src import linear_util as lu
 from jax._src import stages
@@ -983,10 +983,10 @@ def vmap_f(*args, **kwargs):
  axis_size_ = (axis_size if axis_size is not None else
  _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, "vmap"))
  try:
+ axis_data = batching.AxisData(axis_name, axis_size_, spmd_axis_name)
  out_flat = batching.batch(
- flat_fun, axis_name, axis_size_, in_axes_flat,
- lambda: flatten_axes("vmap out_axes", out_tree(), out_axes),
- spmd_axis_name=spmd_axis_name
+ flat_fun, axis_data, in_axes_flat,
+ lambda: flatten_axes("vmap out_axes", out_tree(), out_axes)
  ).call_wrapped(*args_flat)
  except batching.SpecMatchError as e:
  out_axes_flat = flatten_axes("vmap out_axes", out_tree(), out_axes)
@@ -1532,16 +1532,13 @@ def cache_miss(*args, **kwargs):
  is_explicit_global_axis_size=p.is_explicit_global_axis_size,
  )
 
- map_bind_continuation, top_trace, fun_, tracers, params = (
- core.map_bind_with_continuation(pxla.xla_pmap_p, p.flat_fun,
- *p.flat_args, **params))
  execute: Callable | None = None
- if isinstance(top_trace, core.EvalTrace):
- execute = pxla.xla_pmap_impl_lazy(fun_, *tracers, **params)
- out = map_bind_continuation(execute(*tracers))
- else:
- out = map_bind_continuation(
-   pxla.xla_pmap_p.process(top_trace, fun_, tracers, params))
+ with core.take_current_trace() as trace:
+ if isinstance(trace, core.EvalTrace):
+  execute = pxla.xla_pmap_impl_lazy(p.flat_fun, *p.flat_args, **params)
+  out = execute(*p.flat_args)
+ else:
+ out = pxla.xla_pmap_p.bind_with_trace(trace, (p.flat_fun, *p.flat_args), params)
 
  out_tree, out_flat = p.out_tree, out
  out_pytree_def = out_tree()
@@ -1788,7 +1785,7 @@ def linearize(fun: Callable, *primals, has_aux: bool = False
  >>> def f(x): return 3. * jnp.sin(x) + jnp.cos(x / 2.)
  ...
  >>> jax.jvp(f, (2.,), (3.,))
- (Array(3.26819, dtype=float32, weak_type=True), Array(-5.00753, dtype=float32, weak_type=True))
+ (Array(3.2681944, dtype=float32, weak_type=True), Array(-5.007528, dtype=float32, weak_type=True))
  >>> y, f_jvp = jax.linearize(f, 2.)
  >>> print(y)
  3.2681944
@@ -2146,11 +2143,8 @@ def make_jaxpr(
  @wraps(fun)
  @api_boundary
  def make_jaxpr_f(*args, **kwargs):
- with ExitStack() as stack:
- for axis_name, size in axis_env or []:
- stack.enter_context(core.extend_axis_env(axis_name, size, None))
- traced = jit(fun, static_argnums=static_argnums,
- abstracted_axes=abstracted_axes).trace(*args, **kwargs)
+ traced = jit(fun, static_argnums=static_argnums,
+ abstracted_axes=abstracted_axes).trace(*args, **kwargs)
  # `jit` converts tracers in consts to args but that breaks the semantics of
  # `make_jaxpr`. Hence convert the tracers in args back to consts in jaxpr.
  if traced._num_consts:

jax/_src/array.py

@@ -1088,7 +1088,8 @@ def shard_device_array(x, devices, indices, sharding):
  if sharding.is_fully_replicated:
  shards = [x] * len(devices)
  else:
- shards = x._multi_slice(start_indices, limit_indices, removed_dims)
+ with core.set_current_trace(core.EvalTrace()):
+ shards = x._multi_slice(start_indices, limit_indices, removed_dims)
  aval = api_util.shaped_abstractify(x)
  return pxla.batched_device_put(aval, sharding, shards, devices)
 

jax/_src/callback.py

@@ -224,7 +224,8 @@ def pure_callback_lowering(
  ctx, *args, callback: _FlatCallback, sharding: SingleDeviceSharding | None, **params
 ):
  def _callback(*flat_args):
- return tuple(
+ with core.concrete_eval():
+ return tuple(
  pure_callback_impl(
  *flat_args,
  callback=callback,
@@ -429,7 +430,8 @@ def _batch_fun(batched_args):
 
 def io_callback_lowering(ctx, *args, callback, sharding, ordered, **params):
  def _callback(*flat_args):
- return tuple(
+ with core.concrete_eval():
+ return tuple(
  io_callback_impl(
  *flat_args,
  callback=callback,
@@ -511,7 +513,6 @@ def io_callback(
  flat_shape_dtypes, out_tree = tree_util.tree_flatten(result_shape_dtypes)
  flat_result_avals = map(lambda x: core.ShapedArray(x.shape, x.dtype),
  flat_shape_dtypes)
- flat_args = map(core.raise_as_much_as_possible, flat_args)
  out_flat = io_callback_p.bind(
  *flat_args,
  callback=_FlatCallback(callback, in_tree),

jax/_src/checkify.py

@@ -952,7 +952,7 @@ def shard_map_error_check(
 
  if not isinstance(jaxpr, core.ClosedJaxpr):
  jaxpr = core.ClosedJaxpr(jaxpr, ())
- with core.extend_axis_env_nd(mesh.shape.items()):
+ with core.extend_axis_env(mesh.shape.items()):
  # jaxpr to checked_jaxpr
  checked_jaxpr, out_tree, _ = jaxpr_to_checkify_jaxpr(
  jaxpr, enabled_errors, err_tree, *in_avals
@@ -966,7 +966,7 @@ def expand_errors_leading_dim(*xs):
  errs = [lax.expand_dims(e, [0]) for e in errs]
  return *errs, *outs
 
- with core.extend_axis_env_nd(mesh.shape.items()):
+ with core.extend_axis_env(mesh.shape.items()):
  jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
  expand_errors_leading_dim, checked_jaxpr.in_avals
  )

jax/_src/config.py

@@ -858,7 +858,7 @@ class _ThreadLocalExtraJitContext(NamedTuple):
  The initialization, which uses both config.py and core.py is done using
  `_update_thread_local_jit_state` in core.py to prevent circular imports.
  """
- dynamic_trace_state: Any | None = None
+ trace_state: Any | None = None
  axis_env_state: Hashable = ()
  mesh_context_manager: Hashable = ()
  compute_on_context_manager: Hashable = ()