Repository: jax-ml/jax
Pull request #28781 (2025-05-15 21:29:31)

[hijax] prototype hijax pieces
This shouldn't affect any existing behaviors, and it shouldn't regress tracing time noticeably.

The main problem we're trying to solve is `grad`-of-`jit` where the `jit` abstracts a `Box`. In that case, we must pass a reference to the abstracted box to the backward pass. To do that, we make it part of the AD jaxpr. And to avoid polluting all jaxprs, we introduce a new (currently AD-specific) hijax jaxpr, which has more types/primitives than the existing (lojax) jaxprs.

The main implementation ideas:
* each Trace is tagged with a `requires_low: bool`
* each Jaxpr
  * is tagged with an `is_high: bool`, default False but set True while tracing if any hijax primitives are encountered
  * includes an `mut_types: dict[Var, HijaxType]` indicating final types for type-changing mutable hijax types
* each AbstractValue is tagged by a `mutable: bool` which is read to populate `mut_types`
* each Primitive
  * has an `is_high(**params) -> bool` method (depends on params for HOPs)
  * has a `to_lojax(*args, **params)` method taking and returning hijaxtypes-wrapping-lowtracers
* in `Primitive.bind`, we check if `prim.is_high(**params) and trace.requires_low`, and if so we call `prim.to_lojax`

We plan to revise the implementation to look like having a ToLojax trace that sits above the default EvalTrace instance at the top level.

Co-implemented with @dougalm. Design explanation forthcoming.


Patch:
diff --git a/.bazelrc b/.bazelrc
index 367138b2d..53676637c 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -244,10 +244,6 @@ build:ci_linux_aarch64_base --config=clang --verbose_failures=true
 build:ci_linux_aarch64_base --action_env=TF_SYSROOT="/dt10"
 build:ci_linux_aarch64_base --color=yes
 
-# Workaround for https://github.com/numpy/numpy/issues/28843
-# TODO(phawkins): remove this after upgrading to NumPy 2.2.6.
-build:ci_linux_aarch64_base --test_env=OMP_NUM_THREADS=8
-
 build:ci_linux_aarch64 --config=ci_linux_aarch64_base
 build:ci_linux_aarch64 --host_crosstool_top="@ml2014_clang_aarch64_config_aarch64//crosstool:toolchain"
 build:ci_linux_aarch64 --crosstool_top="@ml2014_clang_aarch64_config_aarch64//crosstool:toolchain"
diff --git a/.github/workflows/k8s.yaml b/.github/workflows/k8s.yaml
index 5756b1afb..470a899a1 100644
--- a/.github/workflows/k8s.yaml
+++ b/.github/workflows/k8s.yaml
@@ -1,4 +1,4 @@
-name: Multi-process run using K8s
+name: Distributed run using K8s Jobset
 
 on:
   push:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a66970764..3fcfdb54b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -9,7 +9,7 @@
 
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: cef0300fd0fc4d2a87a85fa2093c6b283ea36f4b  # frozen: v5.0.0
+  rev: 2c9f875913ee60ca25ce70243dc24d5b6415598c  # frozen: v4.6.0
   hooks:
   - id: check-ast
   - id: check-merge-conflict
diff --git a/docs/multi_process.md b/docs/multi_process.md
index f8c2566ca..8ecc51cc2 100644
--- a/docs/multi_process.md
+++ b/docs/multi_process.md
@@ -307,83 +307,6 @@ what it prints:
 
 Woohoo, look at all those TPU cores!
 
-### Kubernetes Example
-
-Running multi-controller JAX on a Kubernetes cluster is almost identical in spirit to the GPU and TPU examples above: every pod runs the same Python program, JAX discovers its peers, and the cluster behaves like one giant machine.
-
-1. **Container image** - start from a JAX-enabled image, e.g. one of the public JAX AI images on Google Artifact Registry ([TPU][google-artifact-tpu] / [GPU][google-artifact-gpu]) or NVIDIA ([NGC][nvidia-ngc] / [JAX-Toolbox][nvidia-jax-toolbox]).
-
-2. **Workload type** - use either a [JobSet][k8s-jobset] or an [indexed Job][k8s-indexed-job]. Each replica corresponds to one JAX process.
-
-3. **Service Account** - JAX needs permission to list the pods that belong to the job so that processes discover their peers. A minimal RBAC setup is provided in [examples/k8s/svc-acct.yaml][rbac-svc-acct].
-
-Below is a [minimal JobSet][minimal-jobset] that launches two replicas. Replace the placeholders - 
-image, GPU count, and any private registry secrets - with values that match your environment.
-
-```yaml
-apiVersion: jobset.x-k8s.io/v1alpha2
-kind: JobSet
-metadata:
-  name: jaxjob
-spec:
-  replicatedJobs:
-  - name: workers
-    template:
-      spec:
-        parallelism: 2
-        completions: 2
-        backoffLimit: 0
-        template:
-          spec:
-            serviceAccountName: jax-job-sa  # kubectl apply -f svc-acct.yaml
-            restartPolicy: Never
-            imagePullSecrets:
-              # https://k8s.io/docs/tasks/configure-pod-container/pull-image-private-registry/
-            - name: null
-            containers:
-            - name: main
-              image: null  # e.g. ghcr.io/nvidia/jax:jax
-              imagePullPolicy: Always
-              resources:
-                limits:
-                  cpu: 1
-                  # https://k8s.io/docs/tasks/manage-gpus/scheduling-gpus/
-                  nvidia.com/gpu: null
-              command: 
-                - python
-              args:
-                - -c
-                - |
-                  import jax
-                  jax.distributed.initialize()
-                  print(jax.devices())
-                  print(jax.local_devices())
-                  assert jax.process_count() > 1
-                  assert len(jax.devices()) > len(jax.local_devices())
-```
-
-Apply the manifest and watch the pods complete:
-
-```bash
-$ kubectl apply -f example.yaml
-$ kubectl get pods -l jobset.sigs.k8s.io/jobset-name=jaxjob
-NAME                       READY   STATUS      RESTARTS   AGE
-jaxjob-workers-0-0-xpx8l   0/1     Completed   0          8m32s
-jaxjob-workers-0-1-ddkq8   0/1     Completed   0          8m32s
-```
-
-When the job finishes, inspect the logs to confirm that every process saw all accelerators:
-
-```bash
-$ kubectl logs -l jobset.sigs.k8s.io/jobset-name=jaxjob
-[CudaDevice(id=0), CudaDevice(id=1)]
-[CudaDevice(id=0)]
-[CudaDevice(id=0), CudaDevice(id=1)]
-[CudaDevice(id=1)]
-```
-
-Every pod should have the same set of global devices and a different set of local devices. At this point, you can replace the inline script with your real JAX program.
-
 Once the processes are set up, we can start building global {class}`jax.Array`s
 and running computations. The remaining Python code examples in this tutorial
 are meant to be run on all processes simultaneously, after running
@@ -657,11 +580,3 @@ assert (np.all(
 [distributed_arrays]: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
 [gpu_machines]: https://cloud.google.com/compute/docs/gpus
 [unified_sharding]: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
-[google-artifact-tpu]: https://console.cloud.google.com/artifacts/docker/cloud-tpu-images/us/jax-ai-image/tpu
-[google-artifact-gpu]: https://console.cloud.google.com/artifacts/docker/deeplearning-images/us-central1/jax-ai-image/gpu
-[nvidia-ngc]: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax
-[nvidia-jax-toolbox]: https://github.com/NVIDIA/JAX-Toolbox
-[k8s-jobset]: https://github.com/kubernetes-sigs/jobset
-[k8s-indexed-job]: https://kubernetes.io/docs/concepts/workloads/controllers/job/#parallel-jobs
-[rbac-svc-acct]: https://github.com/jax-ml/jax/blob/main/examples/k8s/svc-acct.yaml
-[minimal-jobset]: https://github.com/jax-ml/jax/blob/main/examples/k8s/example.yaml
diff --git a/jax/_src/ad_util.py b/jax/_src/ad_util.py
index 8cfd7b214..4e9616e48 100644
--- a/jax/_src/ad_util.py
+++ b/jax/_src/ad_util.py
@@ -31,6 +31,9 @@ T = TypeVar('T')
 map = safe_map
 
 def add_jaxvals(x: ArrayLike, y: ArrayLike) -> Array:
+  ty = core.typeof(x)
+  if hasattr(ty, 'vspace_add'):  # TODO(mattjj,dougalm): revise away hasattr
+    return ty.vspace_add(x, y)
   x, y = core.standard_insert_pvary(x, y)
   return add_jaxvals_p.bind(x, y)
 
@@ -48,6 +51,8 @@ def add_abstract(x, y):
   return x
 
 def zeros_like_aval(aval: core.AbstractValue) -> Array:
+  if hasattr(aval, 'vspace_zero'):  # TODO(mattjj,dougalm): revise away hasattr
+    return aval.vspace_zero()
   return aval_zeros_likers[type(aval)](aval)
 aval_zeros_likers: dict[type, Callable[[Any], Array]] = {}
 
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 154ff5132..33802e494 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -540,17 +540,18 @@ def _check_input_dtype_revderiv(name, holomorphic, allow_int, x):
     if not dtypes.issubdtype(aval.dtype, np.complexfloating):
       raise TypeError(f"{name} with holomorphic=True requires inputs with complex dtype, "
                       f"but got {aval.dtype.name}.")
-  if (dtypes.issubdtype(aval.dtype, dtypes.extended) or
-      dtypes.issubdtype(aval.dtype, np.integer) or
-      dtypes.issubdtype(aval.dtype, np.bool_)):
-    if not allow_int:
-      raise TypeError(f"{name} requires real- or complex-valued inputs (input dtype "
-                      f"that is a sub-dtype of np.inexact), but got {aval.dtype.name}. "
-                      "If you want to use Boolean- or integer-valued inputs, use vjp "
-                      "or set allow_int to True.")
-  elif not dtypes.issubdtype(aval.dtype, np.inexact):
-    raise TypeError(f"{name} requires numerical-valued inputs (input dtype that is a "
-                    f"sub-dtype of np.bool_ or np.number), but got {aval.dtype.name}.")
+  if isinstance(aval, ShapedArray):
+    if (dtypes.issubdtype(aval.dtype, dtypes.extended) or
+        dtypes.issubdtype(aval.dtype, np.integer) or
+        dtypes.issubdtype(aval.dtype, np.bool_)):
+      if not allow_int:
+        raise TypeError(f"{name} requires real- or complex-valued inputs (input dtype "
+                        f"that is a sub-dtype of np.inexact), but got {aval.dtype.name}. "
+                        "If you want to use Boolean- or integer-valued inputs, use vjp "
+                        "or set allow_int to True.")
+    elif not dtypes.issubdtype(aval.dtype, np.inexact):
+      raise TypeError(f"{name} requires numerical-valued inputs (input dtype that is a "
+                      f"sub-dtype of np.bool_ or np.number), but got {aval.dtype.name}.")
 _check_input_dtype_grad = partial(_check_input_dtype_revderiv, "grad")
 
 def _check_output_dtype_revderiv(name, holomorphic, x):
@@ -1873,6 +1874,7 @@ def _jvp(fun: lu.WrappedFun, primals, tangents, has_aux=False):
                     f"structure; primals have tree structure {tree_def} whereas tangents have "
                     f"tree structure {tree_def_2}.")
   for p, t in zip(ps_flat, ts_flat):
+    if not isinstance(core.typeof(p), ShapedArray): continue
     if core.primal_dtype_to_tangent_dtype(_dtype(p)) != _dtype(t):
       raise TypeError("primal and tangent arguments to jax.jvp do not match; "
                       "dtypes must be equal, or in case of int/bool primal dtype "
diff --git a/jax/_src/clusters/k8s_cluster.py b/jax/_src/clusters/k8s_cluster.py
index af1b7c020..b40b39cad 100644
--- a/jax/_src/clusters/k8s_cluster.py
+++ b/jax/_src/clusters/k8s_cluster.py
@@ -107,9 +107,7 @@ class K8sCluster(clusters.ClusterEnv):
           "this job does not have the permission for pod introspection. Please "
           "either grant the default SA permission to read pod info, or create a "
           "dedicated service account with the permission and associated with "
-          "the job. For an example on setting up the service account, see the "
-          "example/k8s directory in the JAX repo. For more details, please refer to "
-          "https://docs.jax.dev/en/latest/multi_process.html#kubernetes-example",
+          "the job. For more details, see <PLACERHOLDER_LINK>.",
           width=80
         ))
       raise RuntimeError('\n'.join(err_msg)) from e
diff --git a/jax/_src/core.py b/jax/_src/core.py
index c730e1c28..12e2763d7 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -88,7 +88,7 @@ DebugInfo = lu.DebugInfo
 
 class Jaxpr:
   __slots__ = ['__weakref__', '_constvars', '_invars', '_outvars', '_eqns',
-               '_effects', '_debug_info']
+               '_effects', '_debug_info', '_is_high', '_mut_types']
 
   _constvars: list[Var]
   _invars: list[Var]
@@ -96,6 +96,8 @@ class Jaxpr:
   _eqns: list[JaxprEqn]
   _effects: Effects
   _debug_info: DebugInfo
+  _is_high: bool
+  _mut_types: dict[Var, Any]
 
   @property
   def constvars(self) -> list[Var]:
@@ -121,6 +123,14 @@ class Jaxpr:
   def debug_info(self) -> DebugInfo:
     return self._debug_info
 
+  @property
+  def is_high(self) -> bool:
+    return self._is_high
+
+  @property
+  def mut_types(self) -> dict[Var, Any]:
+    return self._mut_types
+
   def __init__(self, constvars: Sequence[Var], invars: Sequence[Var],
                outvars: Sequence[Atom], eqns: Sequence[JaxprEqn],
                effects: Effects = no_effects,
@@ -128,6 +138,8 @@ class Jaxpr:
                # compatibility we have to allow calls when the debug_info
                # is missing.
                debug_info: DebugInfo = None,  # type: ignore[annotation-type-mismatch,assignment]
+               is_high: bool = False,
+               mut_types: dict | None = None,
                ):
     """
     Args:
@@ -152,6 +164,8 @@ class Jaxpr:
     # TODO(necula): re-enable these safety checks
     # assert (len(debug_info.arg_names) == len(invars)), (debug_info, invars)
     # assert (len(debug_info.result_paths) == len(outvars)), (debug_info, outvars)
+    self._is_high = is_high
+    self._mut_types = mut_types or {}
 
   def __str__(self):
     return str(self.pretty_print())
@@ -178,6 +192,8 @@ class Jaxpr:
         eqns=kwargs.pop("eqns", self.eqns),
         effects=kwargs.pop("effects", self.effects),
         debug_info=kwargs.pop("debug_info", self.debug_info),
+        is_high=kwargs.pop("is_high", self.is_high),
+        mut_types=kwargs.pop("mut_types", self.mut_types),
     )
     if kwargs:
       raise ValueError(f"Unknown keyword arguments: {kwargs}")
@@ -517,7 +533,7 @@ class Primitive:
     for arg in args:
       if isinstance(arg, Tracer) and not arg._trace.is_valid():
         raise escaped_tracer_error(arg)
-    # TODO: figure out how to handle function arguments
+    # TODO: figure out how to handle function arguments for this assert
     # assert (not config.enable_checks.value or
     #         all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
 
@@ -525,6 +541,10 @@ class Primitive:
     # is called frequently and it's slightly faster to avoid using a context
     # manager object.
     prev_trace = trace_ctx.trace
+
+    if self.is_high(**params) and prev_trace.requires_low:
+      return self.to_lojax(*args, **params)  # type: ignore
+
     trace_ctx.set_trace(eval_trace)
     try:
       return self.bind_with_trace(prev_trace, args, params)
@@ -561,6 +581,9 @@ class Primitive:
   def get_bind_params(self, params):
     return [], params
 
+  def is_high(self, **params) -> bool:
+    return False
+
 
 def _effect_free_abstract_eval(abstract_eval):
   def abstract_eval_(*args, **kwargs):
@@ -627,12 +650,13 @@ def check_avals_context_mesh(avals, prim_name):
 TracerType = TypeVar('TracerType', bound='Tracer')
 
 class Trace(Generic[TracerType]):
-  __slots__ = ("__weakref__", "_invalidated", "_weakref")
+  __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 process_primitive(self, primitive, tracers, params):
     raise NotImplementedError("must override")
@@ -1445,6 +1469,8 @@ def definitely_equal(x, y):
 
 class AbstractValue:
   __slots__: list[str] = []
+  is_high = False
+  mutable = False
 
   def to_tangent_aval(self):
     raise NotImplementedError("must override")
@@ -1948,6 +1974,10 @@ class ShapedArray(UnshapedArray):
     self.sharding = get_sharding(sharding, self.shape)
     self.vma = get_vma(vma, self.sharding.mesh)
 
+  def lower_val(self, val): return [val]
+  def raise_val(self, val): return val
+  def lo_ty(self): return [self]
+
   def update(self, shape=None, dtype=None, weak_type=None, **kwargs):
     if shape is None:
       shape = self.shape
@@ -2777,26 +2807,6 @@ def typematch(t1: AbstractValue, t2: AbstractValue) -> bool:
   else:
     return False
 
-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.vma != a2.vma:
-      mismatches.append('the varying manual axes 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 ''
-
 class JaxprTypeError(TypeError): pass
 
 custom_typechecks: dict[Primitive, Callable] = {}
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 9b28595e1..7b81c4e86 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -917,8 +917,7 @@ def _flatten_bwd(f: Callable,
                "shape/dtypes as the args tuple of the primal function, but at "
                f"output{keystr(kp)} the bwd rule produced an output of "
                f"shape/dtype {a_.str_short()} corresponding "
-               f"to an input of shape/dtype {a.str_short()}"
-               f"{core.aval_mismatch_extra(a, a_)}")
+               f"to an input of shape/dtype {a.str_short()}.")
         raise ValueError(msg)
       results.append(ct)
   return results
diff --git a/jax/_src/interpreters/ad.py b/jax/_src/interpreters/ad.py
index 45705382e..090022c9b 100644
--- a/jax/_src/interpreters/ad.py
+++ b/jax/_src/interpreters/ad.py
@@ -73,6 +73,7 @@ def jvp(fun: lu.WrappedFun, has_aux=False, instantiate=True,
 def jvpfun(f: Callable, instantiate, transform_stack, primals, tangents):
   tag = core.TraceTag()
   tangents = [Zero.from_primal_value(t) if not isinstance(t, Zero)
+              and isinstance(core.typeof(t), core.ShapedArray)
               and dtype(t) == float0 else t for t in tangents]
   ctx = (source_info_util.transform_name_stack('jvp') if transform_stack
          else contextlib.nullcontext())
@@ -475,6 +476,7 @@ class JVPTrace(Trace):
     super().__init__()
     self.tag = tag
     self.parent_trace = parent_trace
+    self.requires_low = False
 
   def to_primal_tangent_pair(self, val):
     if isinstance(val, JVPTracer) and val._trace.tag is self.tag:
@@ -606,7 +608,8 @@ class JVPTrace(Trace):
     return map(partial(maybe_jvp_tracer, self), ps_out, ts_out)
 
 def maybe_jvp_tracer(trace, primal, tangent):
-  if type(tangent) is Zero or dtype(tangent) == float0:
+  if (type(tangent) is Zero or
+      core.typeof(tangent) is core.ShapedArray and dtype(tangent) == float0):
     return primal
   else:
     return JVPTracer(trace, primal, tangent)
@@ -641,6 +644,7 @@ def _primal_tangent_shapes_match(primal, tangent):
   if type(tangent) is not Zero:
     primal_aval = get_aval(primal).strip_weak_type()
     tangent_aval = get_aval(tangent).strip_weak_type()
+    if not isinstance(primal_aval, core.ShapedArray): return  # TODO(mattjj,dougalm)
     assert core.definitely_equal_shape(primal_aval.shape, tangent_aval.shape), (primal_aval.shape, tangent_aval.shape)
     expected_tangent_dtype = core.primal_dtype_to_tangent_dtype(primal_aval.dtype)
     assert expected_tangent_dtype == tangent_aval.dtype, (expected_tangent_dtype, tangent_aval.dtype)
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 5866b0c5f..9e875f43d 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -155,6 +155,7 @@ class JaxprTrace(Trace['JaxprTracer']):
     self.name_stack = name_stack
     self.tag = tag
     self.parent_trace = parent_trace
+    self.requires_low = False
 
   def to_jaxpr_tracer(self, x):
     if isinstance(x, JaxprTracer) and x._trace.tag is self.tag:
@@ -899,9 +900,8 @@ def convert_envvars_to_constvars(jaxpr: Jaxpr, num_env_vars: int) -> Jaxpr:
     raise NotImplementedError
   config.enable_checks.value and core.check_jaxpr(jaxpr)
   env_vars, invars = split_list(jaxpr.invars, [num_env_vars])
-  converted_jaxpr = Jaxpr(constvars=jaxpr.constvars + env_vars,
-                          invars=invars, outvars=jaxpr.outvars, eqns=jaxpr.eqns,
-                          effects=jaxpr.effects, debug_info=jaxpr.debug_info)
+  converted_jaxpr = jaxpr.replace(constvars=jaxpr.constvars + env_vars,
+                                  invars=invars)
   config.enable_checks.value and core.check_jaxpr(converted_jaxpr)
   return converted_jaxpr
 
@@ -1173,6 +1173,7 @@ def _partial_eval_jaxpr_custom_cached(
   out_unknowns = map(op.or_, out_unknowns, ensure_out_unknowns)
   out_inst     = map(op.or_, out_inst,     ensure_out_inst)
 
+
   ins_known, _ = partition_list(in_unknowns, jaxpr.invars)
   outs_known, _ = partition_list(out_unknowns, jaxpr.outvars)
   ref_res_is_input = [r in ins_known for r in residual_refs]
@@ -1181,8 +1182,14 @@ def _partial_eval_jaxpr_custom_cached(
   known_outvars = [*outs_known, *residuals]
   known_effects = make_jaxpr_effects(jaxpr.constvars, ins_known_and_ref_res,
                                      known_outvars, known_eqns)
-  jaxpr_known = Jaxpr(jaxpr.constvars, ins_known_and_ref_res, known_outvars,
-                      known_eqns, known_effects, jaxpr.debug_info)
+  known_mut, staged_mut, ins_known_ = {}, {}, set(ins_known)  # type: ignore
+  for v, t in jaxpr.mut_types.items():
+    [staged_mut, known_mut][v in ins_known_][v] = t
+
+  # TODO(mattjj,necula): debug info should be updated here
+  jaxpr_known = jaxpr.replace(
+      invars=ins_known_and_ref_res, outvars=known_outvars,
+      eqns=known_eqns, effects=known_effects, mut_types=known_mut)
   config.enable_checks.value and core.check_jaxpr(jaxpr_known)
 
   _, ins_staged = partition_list(in_inst, jaxpr.invars)
@@ -1190,9 +1197,10 @@ def _partial_eval_jaxpr_custom_cached(
   staged_invars = [*residuals, *non_input_res_refs, *ins_staged]
   staged_effects = make_jaxpr_effects(jaxpr.constvars, staged_invars,
                                       outs_staged, staged_eqns)
-  jaxpr_staged = Jaxpr(jaxpr.constvars, staged_invars,
-                       outs_staged, staged_eqns, staged_effects,
-                       jaxpr.debug_info)
+  # TODO(mattjj,necula): debug info should be updated here
+  jaxpr_staged = jaxpr.replace(
+      invars=staged_invars, outvars=outs_staged, eqns=staged_eqns,
+      effects=staged_effects, mut_types=staged_mut)
   config.enable_checks.value and core.check_jaxpr(jaxpr_staged)
 
   return (jaxpr_known, jaxpr_staged, out_unknowns, out_inst, len(residuals),
@@ -1483,7 +1491,8 @@ def _dce_jaxpr(jaxpr: Jaxpr, used_outputs: tuple[bool, ...],
       jaxpr.debug_info.traced_for, jaxpr.debug_info.func_src_info,
       jaxpr.debug_info.filter_arg_names(used_inputs),
       jaxpr.debug_info.filter_result_paths(used_outputs))
-  new_jaxpr = Jaxpr(jaxpr.constvars, invars, outvars, eqns, jaxpr_effects, dbg)
+  new_jaxpr = jaxpr.replace(invars=invars, outvars=outvars, eqns=eqns,
+                            effects=jaxpr_effects, debug_info=dbg)
   config.enable_checks.value and core.check_jaxpr(new_jaxpr)
 
   return new_jaxpr, used_inputs
@@ -1561,9 +1570,8 @@ def _move_binders_to_front(closed_jaxpr: ClosedJaxpr, to_move: tuple[bool, ...]
   new_invars = _move_to_front(invars, to_move)
   new_effs = _renumber_effects(
       (*constvars, *new_invars), (*constvars, *invars), closed_jaxpr.jaxpr.effects)
-  new_jaxpr = Jaxpr(constvars, new_invars, closed_jaxpr.jaxpr.outvars,
-                    closed_jaxpr.jaxpr.eqns, new_effs,
-                    closed_jaxpr.jaxpr.debug_info)
+  new_jaxpr = closed_jaxpr.jaxpr.replace(
+      constvars=constvars, invars=new_invars, effects=new_effs)
   new_closed_jaxpr = core.ClosedJaxpr(new_jaxpr, closed_jaxpr.consts)
   return new_closed_jaxpr
 
@@ -1704,6 +1712,7 @@ class JaxprStackFrame:
   attrs_inits: list
   attrs_vars: list[Var]
   debug_info: core.DebugInfo
+  is_high: bool
 
   def __init__(self, debug_info: core.DebugInfo):
     self.gensym = core.gensym()
@@ -1718,6 +1727,7 @@ class JaxprStackFrame:
     self.attrs_inits = []
     self.attrs_vars = []
     self.debug_info = debug_info
+    self.is_high = False
 
   def add_eqn(self, eqn: core.JaxprEqn):
     self.eqns.append(eqn)
@@ -1743,8 +1753,9 @@ class JaxprStackFrame:
     outvars = state_outvars + explicit_outvars
     constvars, constvals = unzip2(self.constvar_to_val.items())
     jaxpr_effects = make_jaxpr_effects(constvars, self.invars, explicit_outvars, self.eqns)
+    mut_types = {v: v.aval for v in invars if v.aval.mutable} if self.is_high else {}
     jaxpr = Jaxpr(constvars, invars, outvars, self.eqns, jaxpr_effects,
-                  debug_info)
+                  debug_info, self.is_high, mut_types)
     jaxpr, constvals = _drop_unused_vars(jaxpr, constvals)
     init_trees = [tree_structure(init_val) for init_val in self.attrs_inits]
     return jaxpr, list(constvals), zip(init_trees, end_trees, self.attrs_tracked)
@@ -1831,8 +1842,9 @@ def _drop_unused_vars(
 class DynamicJaxprTrace(core.Trace):
   __slots__ = ("frame", "tag")
 
-  def __init__(self, debug_info: core.DebugInfo):
+  def __init__(self, debug_info: core.DebugInfo, lower=False):
     super().__init__()
+    self.requires_low = lower
     self.frame = JaxprStackFrame(debug_info)
 
   def invalidate(self):
@@ -2193,10 +2205,11 @@ def trace_to_jaxpr_dynamic(
     in_avals: Sequence[AbstractValue],
     *,
     keep_inputs: list[bool] | None = None,
+    lower: bool = False,
 ) -> tuple[Jaxpr, list[AbstractValue], list[Any],
            list[tuple[PyTreeDef, PyTreeDef, tuple[Any, str, AttrKind]]]]:
   keep_inputs = [True] * len(in_avals) if keep_inputs is None else keep_inputs
-  trace = DynamicJaxprTrace(fun.debug_info)
+  trace = DynamicJaxprTrace(fun.debug_info, lower=lower)
   with core.ensure_no_leaks(trace), source_info_util.reset_name_stack():
     source_info = source_info_util.current()
     in_tracers = _input_type_to_tracers(
@@ -2418,8 +2431,7 @@ def _add_implicit_outputs(jaxpr: Jaxpr) -> tuple[Jaxpr, OutputType]:
   kept_outs = [False] * len(impl_outvars) + [True] * len(expl_outvars)
   out_type = tuple(zip(out_avals, kept_outs))
 
-  new_jaxpr = Jaxpr(jaxpr.constvars, jaxpr.invars, outvars, jaxpr.eqns,
-                    jaxpr.effects, jaxpr.debug_info)
+  new_jaxpr = jaxpr.replace(outvars=outvars)
   config.enable_checks.value and core.check_jaxpr(jaxpr)
   return new_jaxpr, out_type
 
@@ -2663,3 +2675,25 @@ def _linearize_of_pmap_hack(f: lu.WrappedFun, jaxpr, consts) -> tuple[Jaxpr, lis
     _, jaxpr = f.f.closure
     return convert_constvars_jaxpr(jaxpr), []
   return jaxpr, consts
+
+
+@weakref_lru_cache
+def lower_jaxpr(hi_jaxpr):
+  in_avals = [lo_ty for t in hi_jaxpr.in_avals for lo_ty in t.lo_ty()]
+  f = lu.wrap_init(partial(lower_traceable, hi_jaxpr),
+                   debug_info=hi_jaxpr.jaxpr.debug_info)
+  lo_jaxpr, _, consts, () = trace_to_jaxpr_dynamic(f, in_avals, lower=True)
+  return core.ClosedJaxpr(lo_jaxpr, consts)
+
+def lower_traceable(jaxpr, *lo_args):
+  lo_args_ = iter(lo_args)
+  hi_args = [t.raise_val(*it.islice(lo_args_, len(t.lo_ty())))
+             for t in jaxpr.in_avals]
+  assert (problem := next(lo_args_, None)) is None
+  hi_outs = core.jaxpr_as_fun(jaxpr)(*hi_args)
+  in_idx = {v: i for i, v in enumerate(jaxpr.jaxpr.invars)}
+  mut_outs = [lo_val for v, ty in jaxpr.jaxpr.mut_types.items()
+              for lo_val in ty.get(hi_args[in_idx[v]])]
+  lo_outs = [lo_val for t, hi_val in zip(jaxpr.out_avals, hi_outs)
+             for lo_val in t.lower_val(hi_val)]
+  return mut_outs + lo_outs
diff --git a/jax/_src/lax/control_flow/common.py b/jax/_src/lax/control_flow/common.py
index b75cbf6ac..87dbcd8d3 100644
--- a/jax/_src/lax/control_flow/common.py
+++ b/jax/_src/lax/control_flow/common.py
@@ -260,3 +260,23 @@ def _show_diff(array1, array2):
 def _avals_short(avals):
   to_str = lambda aval: getattr(aval, 'str_short', partial(str, aval))()
   return ' '.join(map(to_str, avals))
+
+def _aval_mismatch_extra(a1: core.AbstractValue, a2: core.AbstractValue) -> str:
+  assert not core.typematch(a1, a2)
+  if isinstance(a1, core.ShapedArray) and isinstance(a2, core.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.vma != a2.vma:
+      mismatches.append('the varying manual axes 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 ''
diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
index 741636c47..d87598992 100644
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -53,8 +53,8 @@ from jax._src.lib.mlir.dialects import hlo
 import numpy as np
 
 from jax._src.lax.control_flow.common import (
-    _avals_short, _typecheck_param, _initial_style_jaxprs_with_common_consts,
-    _make_closed_jaxpr, _prune_zeros)
+    _avals_short, _typecheck_param, _aval_mismatch_extra,
+    _initial_style_jaxprs_with_common_consts, _make_closed_jaxpr, _prune_zeros)
 
 map, unsafe_map = safe_map, map
 
@@ -351,7 +351,7 @@ def _check_branch_outputs(
   if not all(map(core.typematch, out_avals1, out_avals2)):
     diffs = [f'the output of {name1}{component(p)} has type {a1.str_short()}'
              f' but the corresponding output of {name2} has type '
-             f'{a2.str_short()}{core.aval_mismatch_extra(a1, a2)}'
+             f'{a2.str_short()}{_aval_mismatch_extra(a1, a2)}'
              for p, a1, a2 in zip(paths, out_avals1, out_avals2)
              if not core.typematch(a1, a2)]
     if len(diffs) == 0:
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index 7efe3294f..c85a23b6b 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -51,7 +51,7 @@ from jax._src.lax import windowed_reductions
 from jax._src.lax.control_flow.common import (
     _avals_short, _initial_style_jaxpr,
     _initial_style_jaxpr_attrs, _make_closed_jaxpr_attrs, _prune_zeros,
-    _typecheck_param)
+    _typecheck_param, _aval_mismatch_extra)
 from jax._src.lax.other import logaddexp
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo
@@ -478,7 +478,7 @@ def _check_carry_type(name, body_fun, in_carry, out_carry_tree, out_avals):
   if not all(_map(core.typematch, in_avals, out_avals)):
     diffs = [f'{component(path)} has type {in_aval.str_short()}'
              ' but the corresponding output carry component has type '
-             f'{out_aval.str_short()}{core.aval_mismatch_extra(in_aval, out_aval)}'
+             f'{out_aval.str_short()}{_aval_mismatch_extra(in_aval, out_aval)}'
              for path, in_aval, out_aval in zip(paths, in_avals, out_avals)
              if not core.typematch(in_aval, out_aval)]
 
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index ca77b659a..5edd74fe7 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -20,7 +20,7 @@ import contextlib
 import dataclasses
 from functools import partial
 import inspect
-import itertools
+import itertools as it
 import logging
 import weakref
 from typing import NamedTuple, Any, Union, cast
@@ -188,7 +188,8 @@ def _python_pjit_helper(fun: Callable, jit_info: PjitInfo, *args, **kwargs):
     args_flat = [*init_states, *args_flat]
 
   try:
-    if core.trace_state_clean() and not config.debug_key_reuse.value:
+    if (core.trace_state_clean() and not config.debug_key_reuse.value
+        and not p.params['jaxpr'].jaxpr.is_high):
       args_flat = map(core.full_lower, args_flat)
       core.check_eval_args(args_flat)
       out_flat, compiled, profiler = _pjit_call_impl_python(*args_flat, **p.params)
@@ -1592,6 +1593,36 @@ pjit_p = core.Primitive("pjit")
 pjit_p.multiple_results = True
 pjit_p.skip_canonicalization = True
 
+def _is_high(jaxpr, **_) -> bool:
+  return jaxpr.jaxpr.is_high
+pjit_p.is_high = _is_high  # type: ignore
+
+def _to_lojax(*hi_args, jaxpr, out_shardings, out_layouts, **params):
+  num_mut = [len(ty.lo_ty()) for ty in jaxpr.jaxpr.mut_types.values()]
+  out_shardings = (UNSPECIFIED,) * sum(num_mut) + out_shardings
+  out_layouts = (None,) * sum(num_mut) + out_layouts
+
+  lo_args = [lo_val for t, hi_val in zip(jaxpr.in_avals, hi_args)
+             for lo_val in t.lower_val(hi_val)]
+  lo_jaxpr = pe.lower_jaxpr(jaxpr)
+  all_outs = pjit_p.bind(*lo_args, jaxpr=lo_jaxpr, out_shardings=out_shardings,
+                         out_layouts=out_layouts, **params)
+  out_mut, lo_outs = split_list(all_outs, [sum(num_mut)])
+
+  out_mut_ = iter(out_mut)
+  in_idx = {v: i for i, v in enumerate(jaxpr.jaxpr.invars)}
+  for var, ty in jaxpr.jaxpr.mut_types.items():
+    ty.set(hi_args[in_idx[var]], *it.islice(out_mut_, len(ty.lo_ty())))
+  assert next(out_mut_, None) is None
+
+  lo_outs_ = iter(lo_outs)
+  hi_outs = [t.raise_val(*it.islice(lo_outs_, len(t.lo_ty())))
+             for t in jaxpr.out_avals]
+  assert next(lo_outs_, None) is None
+
+  return hi_outs
+pjit_p.to_lojax = _to_lojax
+
 def _resolve_in_layouts(args, jit_in_layouts, resolved_in_shardings, in_avals):
   # If device or backend is set, return the default layout. This is because you
   # can pass arrays on cpu (with untiled layouts) to jit with backend='tpu'
@@ -3233,7 +3264,7 @@ def _flatten_boxes(dbg, args, kwargs):
     return args, kwargs, []
   box_data = []
   id_first_occurrences = {}
-  idxs = itertools.count()
+  idxs = it.count()
   def visit(x):
     i = next(idxs)
     if (isinstance(x, (Box, List)) and
diff --git a/jax/numpy/__init__.pyi b/jax/numpy/__init__.pyi
index 4db407861..c52ce2628 100644
--- a/jax/numpy/__init__.pyi
+++ b/jax/numpy/__init__.pyi
@@ -253,8 +253,7 @@ def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int, ...]: ...
 def broadcast_shapes(*shapes: Sequence[int | _core.Tracer]
                      ) -> tuple[int | _core.Tracer, ...]: ...
 
-def broadcast_to(array: ArrayLike, shape: DimSize | Shape, *,
-                 out_sharding: NamedSharding | P | None = None) -> Array: ...
+def broadcast_to(array: ArrayLike, shape: DimSize | Shape) -> Array: ...
 c_: _CClass
 can_cast = _np.can_cast
 def cbrt(x: ArrayLike, /) -> Array: ...
@@ -268,7 +267,6 @@ def clip(
     /,
     min: ArrayLike | None = ...,
     max: ArrayLike | None = ...,
-    *,
     a: ArrayLike | DeprecatedArg | None = ...,
     a_min: ArrayLike | DeprecatedArg | None = ...,
     a_max: ArrayLike | DeprecatedArg | None = ...
@@ -280,7 +278,7 @@ complex128: Any
 complex64: Any
 complex_: Any
 complexfloating = _np.complexfloating
-def compress(condition: ArrayLike, a: ArrayLike, axis: int | None = ..., *,
+def compress(condition: ArrayLike, a: ArrayLike, axis: int | None = ...,
              size: int | None = ..., fill_value: ArrayLike = ..., out: None = ...) -> Array: ...
 def concat(arrays: Sequence[ArrayLike], /, *, axis: int | None = 0) -> Array: ...
 def concatenate(
@@ -316,9 +314,9 @@ def cross(
     axis: int | None = ...,
 ) -> Array: ...
 csingle: Any
-def cumprod(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike | None = ...,
+def cumprod(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike = ...,
             out: None = ...) -> Array: ...
-def cumsum(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike | None = ...,
+def cumsum(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike = ...,
            out: None = ...) -> Array: ...
 def cumulative_prod(x: ArrayLike, /, *, axis: int | None = ...,
                     dtype: DTypeLike | None = ...,
@@ -373,6 +371,7 @@ def einsum(
     optimize: str | builtins.bool | list[tuple[int, ...]] = ...,
     precision: PrecisionLike = ...,
     preferred_element_type: DTypeLike | None = ...,
+    _use_xeinsum: builtins.bool = False,
     _dot_general: Callable[..., Array] = ...,
     out_sharding: NamedSharding | P | None = ...,
 ) -> Array: ...
@@ -386,6 +385,7 @@ def einsum(
     optimize: str | builtins.bool | list[tuple[int, ...]] = ...,
     precision: PrecisionLike = ...,
     preferred_element_type: DTypeLike | None = ...,
+    _use_xeinsum: builtins.bool = False,
     _dot_general: Callable[..., Array] = ...,
     out_sharding: NamedSharding | P | None = ...,
 ) -> Array: ...
@@ -397,6 +397,7 @@ def einsum(
     optimize: str | builtins.bool | list[tuple[int, ...]] = ...,
     precision: PrecisionLike = ...,
     preferred_element_type: DTypeLike | None = ...,
+    _use_xeinsum: builtins.bool = ...,
     _dot_general: Callable[..., Array] = ...,
     out_sharding: NamedSharding | P | None = ...,
 ) -> Array: ...
@@ -421,7 +422,7 @@ def einsum_path(
     optimize: str | builtins.bool | list[tuple[int, ...]] =  ...,
 ) -> tuple[list[tuple[int, ...]], Any]: ...
 
-def empty(shape: Any, dtype: DTypeLike | None = ..., *,
+def empty(shape: Any, dtype: DTypeLike | None = ...,
           device: _Device | _Sharding | None = ...) -> Array: ...
 def empty_like(prototype: ArrayLike | DuckTypedArray,
                dtype: DTypeLike | None = ...,
@@ -578,17 +579,17 @@ def intersect1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: builtins.bool = .
 def invert(x: ArrayLike, /) -> Array: ...
 def isclose(a: ArrayLike, b: ArrayLike, rtol: ArrayLike = ...,
             atol: ArrayLike = ..., equal_nan: builtins.bool = ...) -> Array: ...
-def iscomplex(x: ArrayLike) -> Array: ...
+def iscomplex(m: ArrayLike) -> Array: ...
 def iscomplexobj(x: Any) -> builtins.bool: ...
 def isdtype(dtype: DTypeLike, kind: DType | str | tuple[DType | str, ...]) -> builtins.bool: ...
 def isfinite(x: ArrayLike, /) -> Array: ...
-def isin(element: ArrayLike, test_elements: ArrayLike, assume_unique: builtins.bool = ...,
-         invert: builtins.bool = ..., *, method: str = ...) -> Array: ...
+def isin(element: ArrayLike, test_elements: ArrayLike,
+         assume_unique: builtins.bool = ..., invert: builtins.bool = ..., method: str = ...) -> Array: ...
 def isinf(x: ArrayLike, /) -> Array: ...
 def isnan(x: ArrayLike, /) -> Array: ...
 def isneginf(x: ArrayLike, /) -> Array: ...
 def isposinf(x: ArrayLike, /) -> Array: ...
-def isreal(x: ArrayLike) -> Array: ...
+def isreal(m: ArrayLike) -> Array: ...
 def isrealobj(x: Any) -> builtins.bool: ...
 def isscalar(element: Any) -> builtins.bool: ...
 def issubdtype(arg1: DTypeLike, arg2: DTypeLike) -> builtins.bool: ...
@@ -643,7 +644,7 @@ def logspace(start: ArrayLike, stop: ArrayLike, num: int = ...,
              endpoint: builtins.bool = ..., base: ArrayLike = ...,
              dtype: DTypeLike | None = ..., axis: int = ...) -> Array: ...
 def mask_indices(
-    n: int, mask_func: Callable, k: int = ..., *, size: int | None = ...
+    n: int, mask_func: Callable, k: int = ...
 ) -> tuple[Array, ...]: ...
 def matmul(
     a: ArrayLike, b: ArrayLike, *, precision: PrecisionLike = ...,
@@ -654,7 +655,7 @@ def max(a: ArrayLike, axis: _Axis = ..., out: None = ...,
         keepdims: builtins.bool = ..., initial: ArrayLike | None = ...,
         where: ArrayLike | None = ...) -> Array: ...
 def maximum(x: ArrayLike, y: ArrayLike, /) -> Array: ...
-def mean(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def mean(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
          out: None = ..., keepdims: builtins.bool = ..., *,
          where: ArrayLike | None = ...) -> Array: ...
 def median(a: ArrayLike, axis: int | tuple[int, ...] | None = ...,
@@ -688,14 +689,14 @@ def nanargmin(
     out: None = ...,
     keepdims: builtins.bool | None = ...,
 ) -> Array: ...
-def nancumprod(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike | None = ...,
+def nancumprod(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike = ...,
                out: None = ...) -> Array: ...
-def nancumsum(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike | None = ...,
+def nancumsum(a: ArrayLike, axis: int | None = ..., dtype: DTypeLike = ...,
                out: None = ...) -> Array: ...
 def nanmax(a: ArrayLike, axis: _Axis = ..., out: None = ...,
            keepdims: builtins.bool = ..., initial: ArrayLike | None = ...,
            where: ArrayLike | None = ...) -> Array: ...
-def nanmean(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def nanmean(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
             out: None = ...,
             keepdims: builtins.bool = ...,
             where: ArrayLike | None = ...) -> Array: ...
@@ -709,21 +710,21 @@ def nanpercentile(a: ArrayLike, q: ArrayLike,
                   axis: int | tuple[int, ...] | None = ...,
                   out: None = ..., overwrite_input: builtins.bool = ..., method: str = ...,
                   keepdims: builtins.bool = ..., *, interpolation: DeprecatedArg | str = ...) -> Array: ...
-def nanprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def nanprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
             out: None = ...,
             keepdims: builtins.bool = ..., initial: ArrayLike | None = ...,
             where: ArrayLike | None = ...) -> Array: ...
 def nanquantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = ...,
                 out: None = ..., overwrite_input: builtins.bool = ..., method: str = ...,
                 keepdims: builtins.bool = ..., *, interpolation: DeprecatedArg | str = ...) -> Array: ...
-def nanstd(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
-           out: None = ..., ddof: int = ..., keepdims: builtins.bool = ...,
+def nanstd(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ..., out: None = ...,
+           ddof: int = ..., keepdims: builtins.bool = ...,
            where: ArrayLike | None = ...) -> Array: ...
-def nansum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def nansum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ..., keepdims: builtins.bool = ...,
            initial: ArrayLike | None = ...,
            where: ArrayLike | None = ...) -> Array: ...
-def nanvar(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def nanvar(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ...,
            ddof: int = 0, keepdims: builtins.bool = False,
            where: ArrayLike | None = ...) -> Array: ...
@@ -739,7 +740,7 @@ def not_equal(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 number = _np.number
 object_ = _np.object_
 ogrid: _Ogrid
-def ones(shape: Any, dtype: DTypeLike | None = ..., *,
+def ones(shape: Any, dtype: DTypeLike | None = ...,
          device: _Device | _Sharding | None = ...) -> Array: ...
 def ones_like(a: ArrayLike | DuckTypedArray,
               dtype: DTypeLike | None = ...,
@@ -781,7 +782,7 @@ def positive(x: ArrayLike, /) -> Array: ...
 def pow(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def power(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 printoptions = _np.printoptions
-def prod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def prod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
          out: None = ..., keepdims: builtins.bool = ...,
          initial: ArrayLike | None = ..., where: ArrayLike | None = ...,
          promote_integers: builtins.bool = ...) -> Array: ...
@@ -804,6 +805,7 @@ def ravel_multi_index(multi_index: Sequence[ArrayLike], dims: Sequence[int],
                       mode: str = ..., order: str = ...) -> Array: ...
 def real(x: ArrayLike, /) -> Array: ...
 def reciprocal(x: ArrayLike, /) -> Array: ...
+register_jax_array_methods: Any
 def remainder(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = ..., *,
            total_repeat_length: int | None = ...,
@@ -842,8 +844,7 @@ def setdiff1d(
     size: int | None = ...,
     fill_value: ArrayLike | None = ...,
 ) -> Array: ...
-def setxor1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: builtins.bool = ..., *,
-             size: int | None = ..., fill_value: ArrayLike | None = ...) -> Array: ...
+def setxor1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: builtins.bool = ...) -> Array: ...
 def shape(a: ArrayLike | SupportsShape) -> tuple[int, ...]: ...
 def sign(x: ArrayLike, /) -> Array: ...
 def signbit(x: ArrayLike, /) -> Array: ...
@@ -881,14 +882,14 @@ def stack(
     out: None = ...,
     dtype: DTypeLike | None = ...,
 ) -> Array: ...
-def std(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def std(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
         out: None = ..., ddof: int = ..., keepdims: builtins.bool = ..., *,
         where: ArrayLike | None = ..., correction: int | float | None = ...) -> Array: ...
 subtract: BinaryUfunc
 def sum(
     a: ArrayLike,
     axis: _Axis = ...,
-    dtype: DTypeLike | None = ...,
+    dtype: DTypeLike = ...,
     out: None = ...,
     keepdims: builtins.bool = ...,
     initial: ArrayLike | None = ...,
@@ -926,7 +927,7 @@ def transpose(a: ArrayLike, axes: Sequence[int] | None = ...) -> Array: ...
 def trapezoid(y: ArrayLike, x: ArrayLike | None = None, dx: ArrayLike = ...,
               axis: int = ...) -> Array: ...
 def tri(
-    N: int, M: int | None = ..., k: int = ..., dtype: DTypeLike | None = ...
+    N: int, M: int | None = ..., k: int = ..., dtype: DTypeLike = ...
 ) -> Array: ...
 def tril(m: ArrayLike, k: int = ...) -> Array: ...
 def tril_indices(
@@ -969,7 +970,7 @@ class _UniqueInverseResult(NamedTuple):
 def unique(ar: ArrayLike, return_index: builtins.bool = ..., return_inverse: builtins.bool = ...,
            return_counts: builtins.bool = ..., axis: int | None = ...,
            *, equal_nan: builtins.bool = ..., size: int | None = ...,
-           fill_value: ArrayLike | None = ..., sorted: bool = ...,
+           fill_value: ArrayLike | None = ...
 ): ...
 def unique_all(x: ArrayLike, /, *, size: int | None = ...,
                fill_value: ArrayLike | None = ...) -> _UniqueAllResult: ...
@@ -993,7 +994,7 @@ def unwrap(p: ArrayLike, discont: ArrayLike | None = ...,
 def vander(
     x: ArrayLike, N: int | None = ..., increasing: builtins.bool = ...
 ) -> Array: ...
-def var(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike | None = ...,
+def var(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
         out: None = ..., ddof: int = ..., keepdims: builtins.bool = ..., *,
         where: ArrayLike | None = ..., correction: int | float | None = ...) -> Array: ...
 def vdot(
@@ -1028,7 +1029,7 @@ def where(condition: ArrayLike, x: ArrayLike | None = ...,
           fill_value: None | ArrayLike | tuple[ArrayLike, ...] = ...
           ) -> Array | tuple[Array, ...]: ...
 
-def zeros(shape: Any, dtype: DTypeLike | None = ..., *,
+def zeros(shape: Any, dtype: DTypeLike | None = ...,
           device: _Device | _Sharding | None = ...) -> Array: ...
 def zeros_like(a: ArrayLike | DuckTypedArray,
                dtype: DTypeLike | None = ...,


Test patch:
diff --git a/tests/attrs_test.py b/tests/attrs_test.py
index 60a3753a7..b6cef7fec 100644
--- a/tests/attrs_test.py
+++ b/tests/attrs_test.py
@@ -15,7 +15,9 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import partial
 import itertools as it
+import unittest
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -25,6 +27,10 @@ import jax
 import jax.numpy as jnp
 
 from jax._src import config
+from jax._src import core
+from jax._src import dtypes
+from jax._src.interpreters import ad
+from jax._src.interpreters import partial_eval as pe
 from jax._src import test_util as jtu
 from jax._src.util import safe_zip, safe_map
 
@@ -1326,5 +1332,292 @@ class ListTest(jtu.JaxTestCase):
       f(b, b)
 
 
+class HiPrimitive(core.Primitive):
+  def __init__(self, name):
+    self.name = name
+    ad.primitive_jvps[self] = self.jvp
+    ad.primitive_transposes[self] = self.transpose
+    pe.custom_staging_rules[self] = self.staging
+
+  def staging(self, trace, *args, **kwargs):
+    trace.frame.is_high = True
+    return trace.default_process_primitive(self, args, kwargs)
+
+  def is_high(self, **params):
+    return True
+
+  def abstract_eval(self, *arg_avals, **params):
+    assert False, "must override"
+
+  def to_lojax(self, *lotypes_wrapped_in_hitypes, **params):
+    assert False, "must override"
+
+  def jvp(self, primals, tangents, **params):
+    assert False, "must override"
+
+  def transpose(self, *args, **params):
+    assert False  # TODO
+
+
+class HijaxTest(jtu.JaxTestCase):
+
+  def test_custom_types_and_primitive(self):
+    if config.enable_x64.value: raise unittest.SkipTest("no x64")
+
+    @dataclass(frozen=True)
+    class MyArray:
+      arr: jax.Array  # always f32
+
+    @dataclass(frozen=True)
+    class MyTy(core.AbstractValue):
+      mutable = False
+
+      def to_tangent_aval(self):
+        return MyTy()
+      def str_short(self, short_dtypes=False):
+        return 'MyTy'
+      def lo_ty(self):
+        return [core.ShapedArray((), jnp.dtype('float32'))]
+      def lower_val(self, hi_val: MyArray) -> list[jax.Array]:
+        return [hi_val.arr]
+      def raise_val(self, val) -> MyArray:
+        return MyArray(val)
+
+      def __eq__(self, other): return isinstance(other, MyTy)
+
+      def vspace_zero(self):
+        return MyArray(jnp.zeros((), 'float32'))
+      def vspace_add(self, x, y):
+        return add(x, y)
+
+      def strip_weak_type(self): return self
+      def normalize(self): return self
+    core.pytype_aval_mappings[MyArray] = lambda _: MyTy()
+
+    class ToMy(HiPrimitive):
+      def is_high(self): return True
+
+      def abstract_eval(_, lo_aval):
+        return MyTy(), set()
+
+      def to_lojax(_, lo):
+        return MyArray(lo)
+
+      def jvp(_, primals, tangents):
+        x, x_dot = *primals, *tangents
+        return to(x), to(x_dot)
+
+      def transpose(self, out_bar, _):
+        return from_(out_bar),
+
+    class FromMy(HiPrimitive):
+      def is_high(self): return True
+
+      def abstract_eval(_, hi_aval):
+        return hi_aval.lo_ty()[0], set()
+
+      def to_lojax(_, hi):
+        return hi.arr
+
+      def jvp(_, primals, tangents):
+        x, x_dot = *primals, *tangents
+        return from_(x), from_(x_dot)
+
+      def transpose(self, out_bar, _):
+        return to(out_bar),
+
+    def to(x): return to_p.bind(x)
+    to_p = ToMy('to_my')
+
+    def from_(x): return from_p.bind(x)
+    from_p = FromMy('from_my')
+
+    def mul(x, y): return mul_p.bind(x, y)
+    def add(x, y): return add_p.bind(x, y)
+
+    class MyMul(HiPrimitive):
+      def is_high(self): return True
+
+      def abstract_eval(_, hi_x, hi_y):
+        if hi_x != hi_y: raise Exception
+        return hi_x, set()
+
+      def to_lojax(_, hi_x, hi_y):
+        return MyArray(hi_x.arr * hi_y.arr)
+
+      def jvp(_, primals, tangents):
+        (x, y), (x_dot, y_dot) = primals, tangents
+        return mul(x, y), add(mul(x, y_dot), mul(x_dot, y))
+
+      def transpose(self, out_bar, x, y):
+        assert ad.is_undefined_primal(x) ^ ad.is_undefined_primal(y)
+        if ad.is_undefined_primal(x):
+          return mul(out_bar, y), None
+        else:
+          return None, mul(x, out_bar)
+
+    class MyAdd(HiPrimitive):
+      def is_high(self): return True
+
+      def abstract_eval(_, hi_x, hi_y):
+        if hi_x != hi_y: raise Exception
+        return hi_x, set()
+
+      def to_lojax(_, hi_x, hi_y):
+        return MyArray(hi_x.arr + hi_y.arr)
+
+      def jvp(_, primals, tangents):
+        assert False  # TODO
+
+      def transpose(self, out_bar, x, y):
+        return out_bar, out_bar
+
+    mul_p = MyMul('my_mul')
+    add_p = MyAdd('my_add')
+
+
+    @jax.jit
+    def f(x):
+      return to(from_(x))
+
+    # test basic to/from jit
+    a = MyArray(jnp.ones(()))
+    b = f(a)  # don't crash
+    self.assertIsInstance(b, MyArray)
+    self.assertAllClose(b.arr, jnp.ones(()))
+
+    # test basic to/from autodiff
+    b, b_dot = jax.jvp(f, (a,), (a,))
+    self.assertIsInstance(b, MyArray)
+    self.assertIsInstance(b_dot, MyArray)
+
+    # test mul jit and backward pass
+
+    @jax.jit
+    def f(x):
+      return mul(x, x)
+
+    b, f_vjp = jax.vjp(f, a)
+    self.assertIn('MyTy', str(f_vjp))
+    a_grad, = f_vjp(b)
+    self.assertIsInstance(a_grad, MyArray)
+    self.assertAllClose(a_grad.arr, 2.0, check_dtypes=False)
+
+  def test_box_autodiff(self):
+    if config.enable_x64.value: raise unittest.SkipTest("no x64")
+    class BoxTy(core.AbstractValue):
+      mutable = True
+
+      def to_tangent_aval(self):
+        # NOTE not really used, for some reason we had to write it anyway
+        return core.ShapedArray((), dtypes.float0)
+
+      def str_short(self, short_dtypes=False):
+        return 'BoxTy'
+
+      def lower_val(self, box):
+        return [box._val]
+
+      def raise_val(self, val):
+        return Box(val)  # we're gonna mutate this
+
+      def lo_ty(self):
+        return [core.ShapedArray((), jnp.dtype('float32'))]
+
+      def get(self, box):
+        return [box._val]
+
+      def set(self, box, val):
+        box._val = val
+
+    class Box:
+      def __init__(self, val):
+        self._val = val
+      ty = BoxTy()
+    core.pytype_aval_mappings[Box] = lambda b: b.ty
+
+
+    class BoxSet(HiPrimitive):
+      multiple_results = True
+      def is_high(self) -> bool: return True
+
+      def abstract_eval(*_, **__):
+        return [], set()
+
+      def to_lojax(_, box, val):
+        box._val = val
+        return []
+
+      def jvp(_, primals, tangents):
+        assert False  # TODO
+
+      def transpose(_, *args):
+        assert False  # TODO
+    box_set_p = BoxSet('box_set')
+
+    class BoxGet(HiPrimitive):
+      def is_high(self) -> bool: return True
+
+      def abstract_eval(*_, **__):
+        return jnp.dtype('float32'), set()
+
+      def to_lojax(_, box):
+        return box._val
+
+      def jvp(_, primals, tangents):
+        assert False  # TODO
+
+      def transpose(_, *args):
+        assert False  # TODO
+    box_get_p = BoxGet('box_get')
+
+    class StashTangents(HiPrimitive):
+      def is_high(self):
+        return True
+
+      def abstract_eval(_, box_aval, x_aval):
+        del box_aval
+        return x_aval, set()
+
+      def to_lojax(_, box, x):
+        assert False  # TODO
+
+      def jvp(_, primals, tangents):
+        box, x = primals
+        _, x_dot = tangents
+        box_set(box, x_dot)
+        return x, x_dot
+
+      def transpose(self, *args):
+        assert False  # TODO
+    stash_tangents_p = StashTangents('stash_tangents')
+
+    def box_set(box, val):
+      box_set_p.bind(box, val)
+
+    def box_get(box):
+      return box_get_p.bind(box)
+
+    def stash_tangents(box, x):
+      return stash_tangents_p.bind(box, x)
+
+    @jax.jit
+    def f(box, x):
+      box_set(box, x)
+
+    box = Box(0.0)
+    f(box, 1.)
+    self.assertAllClose(box_get(box), 1.0, check_dtypes=False)
+
+    @jax.jit
+    def f(box, x):
+      x = stash_tangents(box, x)
+      return x
+
+    box = Box(0.0)
+    jax.jvp(partial(f, box), (3.,), (5.,))
+    self.assertAllClose(box_get(box), 5.0, check_dtypes=False)
+
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())