Repository: jax-ml/jax
Pull request #28804 (2025-05-16 22:22:38)

[hijax] type-changing boxes with pytree contents
Allow hijax-style Boxes to have pytree contents. Since we lower away pytrees before hijax (ie before bind), that means flattening in the box_set traceable. The BoxTy type is still parameterized by the treedef, though.

This was also a test for a hijax type that lowers to multiple lojax types. A better warmup would've been to add hijax tuples.

To handle type-changing, a hijaxpr has a `final_typechange_env: dict[Var, HijaxType]`. The initial types are currently attached to the Box lambda binders.


Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index ab97c7ff1..e49173c3d 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', '_is_high', '_mut_types']
+               '_effects', '_debug_info', '_is_high', '_final_typechange_env']
 
   _constvars: list[Var]
   _invars: list[Var]
@@ -97,7 +97,7 @@ class Jaxpr:
   _effects: Effects
   _debug_info: DebugInfo
   _is_high: bool
-  _mut_types: dict[Var, Any]
+  _final_typechange_env: dict[Var, Any]
 
   @property
   def constvars(self) -> list[Var]:
@@ -128,8 +128,8 @@ class Jaxpr:
     return self._is_high
 
   @property
-  def mut_types(self) -> dict[Var, Any]:
-    return self._mut_types
+  def final_typechange_env(self) -> dict[Var, Any]:
+    return self._final_typechange_env
 
   def __init__(self, constvars: Sequence[Var], invars: Sequence[Var],
                outvars: Sequence[Atom], eqns: Sequence[JaxprEqn],
@@ -139,7 +139,7 @@ class Jaxpr:
                # is missing.
                debug_info: DebugInfo = None,  # type: ignore[annotation-type-mismatch,assignment]
                is_high: bool = False,
-               mut_types: dict | None = None,
+               final_typechange_env: dict | None = None,
                ):
     """
     Args:
@@ -165,7 +165,7 @@ class Jaxpr:
     # 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 {}
+    self._final_typechange_env = final_typechange_env or {}
 
   def __str__(self):
     return str(self.pretty_print())
@@ -193,7 +193,8 @@ class Jaxpr:
         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),
+        final_typechange_env=kwargs.pop("final_typechange_env",
+                                        self.final_typechange_env),
     )
     if kwargs:
       raise ValueError(f"Unknown keyword arguments: {kwargs}")
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 9e875f43d..f77db5443 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -1183,13 +1183,13 @@ def _partial_eval_jaxpr_custom_cached(
   known_effects = make_jaxpr_effects(jaxpr.constvars, ins_known_and_ref_res,
                                      known_outvars, known_eqns)
   known_mut, staged_mut, ins_known_ = {}, {}, set(ins_known)  # type: ignore
-  for v, t in jaxpr.mut_types.items():
+  for v, t in jaxpr.final_typechange_env.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)
+      eqns=known_eqns, effects=known_effects, final_typechange_env=known_mut)
   config.enable_checks.value and core.check_jaxpr(jaxpr_known)
 
   _, ins_staged = partition_list(in_inst, jaxpr.invars)
@@ -1200,7 +1200,7 @@ def _partial_eval_jaxpr_custom_cached(
   # 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)
+      effects=staged_effects, final_typechange_env=staged_mut)
   config.enable_checks.value and core.check_jaxpr(jaxpr_staged)
 
   return (jaxpr_known, jaxpr_staged, out_unknowns, out_inst, len(residuals),
@@ -1713,6 +1713,7 @@ class JaxprStackFrame:
   attrs_vars: list[Var]
   debug_info: core.DebugInfo
   is_high: bool
+  final_typechange_env: dict
 
   def __init__(self, debug_info: core.DebugInfo):
     self.gensym = core.gensym()
@@ -1728,6 +1729,7 @@ class JaxprStackFrame:
     self.attrs_vars = []
     self.debug_info = debug_info
     self.is_high = False
+    self.final_typechange_env = {}
 
   def add_eqn(self, eqn: core.JaxprEqn):
     self.eqns.append(eqn)
@@ -1753,9 +1755,8 @@ 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, self.is_high, mut_types)
+                  debug_info, self.is_high, self.final_typechange_env)
     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)
@@ -1872,6 +1873,8 @@ class DynamicJaxprTrace(core.Trace):
     self.frame.tracers.append(tracer)
     self.frame.tracer_to_var[id(tracer)] = var = self.frame.newvar(aval)
     self.frame.invars.append(var)
+    if aval.mutable:
+      self.frame.final_typechange_env[var] = aval
     return tracer
 
   def new_const(self, c, source_info: SourceInfo):
@@ -2692,7 +2695,7 @@ def lower_traceable(jaxpr, *lo_args):
   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()
+  mut_outs = [lo_val for v, ty in jaxpr.jaxpr.final_typechange_env.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)]
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 5edd74fe7..10e7e697e 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -1597,21 +1597,18 @@ 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
+def _to_lojax( *hi_args, jaxpr, **params):
+  params, num_mutants = _lojax_expand_params(jaxpr, **params)
 
   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)])
+  all_outs = pjit_p.bind(*lo_args, jaxpr=lo_jaxpr, **params)
+  out_mut, lo_outs = split_list(all_outs, [num_mutants])
 
   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():
+  for var, ty in jaxpr.jaxpr.final_typechange_env.items():
     ty.set(hi_args[in_idx[var]], *it.islice(out_mut_, len(ty.lo_ty())))
   assert next(out_mut_, None) is None
 
@@ -1623,6 +1620,31 @@ def _to_lojax(*hi_args, jaxpr, out_shardings, out_layouts, **params):
   return hi_outs
 pjit_p.to_lojax = _to_lojax
 
+def _lojax_expand_params(
+    hi_jaxpr, *, donated_invars, in_shardings, in_layouts, out_shardings,
+    out_layouts, **params):
+  # some pjit params match the length of hi_jaxpr.invars/outvars, so when
+  # lowering we must expand them to match their number of lojax types
+  def expand(hi_tys, xs):
+    return tuple(y for hi, x in zip(hi_tys, xs) for y in (x,) * len(hi.lo_ty()))
+  donated_invars = expand(hi_jaxpr.in_avals , donated_invars)
+  in_shardings   = expand(hi_jaxpr.in_avals , in_shardings  )
+  in_layouts     = expand(hi_jaxpr.in_avals , in_layouts    )
+  out_shardings  = expand(hi_jaxpr.out_avals, out_shardings )
+  out_layouts    = expand(hi_jaxpr.out_avals, out_layouts   )
+
+  # also, the lo_jaxpr has pure outputs corresponding to mutable hi_jaxpr types
+  num_mutants = sum(len(hi_ty.lo_ty()) for hi_ty in
+                    hi_jaxpr.jaxpr.final_typechange_env.values())
+  out_shardings = (UNSPECIFIED,) * num_mutants + out_shardings
+  out_layouts = (None,) * num_mutants + out_layouts
+
+  new_params = dict(params, donated_invars=donated_invars,
+                    in_shardings=in_shardings, in_layouts=in_layouts,
+                    out_shardings=out_shardings, out_layouts=out_layouts)
+  return new_params, num_mutants
+
+
 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'


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index c51d40715..35695c75c 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1875,6 +1875,18 @@ jax_multiplatform_test(
     ]),
 )
 
+jax_multiplatform_test(
+    name = "hijax_test",
+    srcs = ["hijax_test.py"],
+    deps = [
+        "//jax:experimental",
+    ] + py_deps([
+        "numpy",
+        "absl/testing",
+    ]),
+)
+
+
 jax_multiplatform_test(
     name = "colocated_python_test",
     srcs = ["colocated_python_test.py"],
diff --git a/tests/attrs_test.py b/tests/attrs_test.py
index b6cef7fec..60a3753a7 100644
--- a/tests/attrs_test.py
+++ b/tests/attrs_test.py
@@ -15,9 +15,7 @@
 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
@@ -27,10 +25,6 @@ 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
 
@@ -1332,292 +1326,5 @@ 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())
diff --git a/tests/hijax_test.py b/tests/hijax_test.py
new file mode 100644
index 000000000..21034d164
--- /dev/null
+++ b/tests/hijax_test.py
@@ -0,0 +1,453 @@
+# Copyright 2024 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 dataclasses import dataclass
+from functools import partial
+import itertools as it
+import unittest
+
+from absl.testing import absltest
+
+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
+
+config.parse_flags_with_absl()
+
+map, unsafe_map = safe_map, map
+zip, unsafe_zip = safe_zip, zip
+
+
+# TODO(mattjj,dougalm): move HiPrimitive, Box, etc out of tests and into library
+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 BoxTy(core.AbstractValue):
+  mutable = True
+
+  def __init__(self, leaf_avals, treedef):
+    self._leaf_avals = leaf_avals  # hijax avals
+    self._treedef = treedef
+
+  # aval interface: hashability and str_short
+  def __hash__(self):
+    return hash((self._leaf_avals, self._treedef))
+
+  def __eq__(self, other):
+    return (isinstance(other, BoxTy) and self._leaf_avals == other._leaf_avals
+            and self._treedef == other._treedef)
+
+  def str_short(self, short_dtypes=False):
+    return 'BoxTy'
+
+  # hijax interface: lower val, raise val, and low type
+  def lo_ty(self):
+    return [lo_aval for hi_aval in self._leaf_avals for lo_aval in hi_aval.lo_ty()]
+
+  def lower_val(self, box):
+    leaf_vals, treedef = jax.tree.flatten(box._val)
+    assert treedef == self._treedef
+    return [lo_val for hi_aval, hi_val in zip(self._leaf_avals, leaf_vals)
+            for lo_val in hi_aval.lower_val(hi_val)]
+
+  def raise_val(self, *lo_vals):
+    lo_vals_ = iter(lo_vals)
+    hi_vals = [hi_ty.raise_val(*it.islice(lo_vals_, len(hi_ty.lo_ty())))
+               for hi_ty in self._leaf_avals]
+    assert next(lo_vals_, None) is None
+    return Box(jax.tree.unflatten(self._treedef, hi_vals))  # will be mutated
+
+  # mutable interface: get/set
+  def get(self, box):
+    leaf_vals, treedef = jax.tree.flatten(box._val)
+    assert treedef == self._treedef
+    return [lo_val for hi_ty, hi_val in zip(self._leaf_avals, leaf_vals)
+            for lo_val in hi_ty.lower_val(hi_val)]
+
+  def set(self, box, *lo_vals):
+    lo_vals_ = iter(lo_vals)
+    hi_vals = [hi_ty.raise_val(*it.islice(lo_vals_, len(hi_ty.lo_ty())))
+               for hi_ty in self._leaf_avals]
+    assert next(lo_vals_, None) is None
+    box._val = jax.tree.unflatten(self._treedef, hi_vals)
+
+  # TODO placeholder thing
+  def to_tangent_aval(self):
+    return core.ShapedArray((), dtypes.float0)  # TODO revise placeholder
+
+class Box:  # noqa: F811
+  def __init__(self, val):
+    self._val = val
+
+  @property
+  def ty(self):
+    leaves, treedef = jax.tree.flatten(self._val)
+    leaf_avals = tuple(map(core.typeof, leaves))
+    return BoxTy(leaf_avals, treedef)
+core.pytype_aval_mappings[Box] = lambda b: b.ty
+
+
+class BoxSet(HiPrimitive):
+  multiple_results = True
+
+  def is_high(self, *, treedef) -> bool: return True
+
+  def staging(self, trace, box, *leaves, treedef):
+    super().staging(trace, box, *leaves, treedef=treedef)
+    avals = tuple(t.aval for t in leaves)
+    trace.frame.final_typechange_env[trace.getvar(box)] = BoxTy(avals, treedef)
+
+  def abstract_eval(self, box_ty, *leaf_avals, treedef):
+    return [], set()  # TODO better typechecking...
+
+  def to_lojax(_, box, *leaves, treedef):
+    box._val = jax.tree.unflatten(treedef, leaves)
+    return []
+
+  def jvp(_, primals, tangents, *, treedef):
+    assert False  # TODO
+
+  def transpose(_, *args, treedef):
+    assert False  # TODO
+box_set_p = BoxSet('box_set')
+
+def box_set(box, val):
+  leaves, treedef = jax.tree.flatten(val)
+  box_set_p.bind(box, *leaves, treedef=treedef)
+
+
+class BoxGet(HiPrimitive):
+  multiple_results = True
+
+  def is_high(self) -> bool: return True
+
+  def abstract_eval(self, box_ty):
+    return box_ty._leaf_avals, set()
+
+  def to_lojax(_, box):
+    return jax.tree.leaves(box._val)
+
+  def jvp(_, primals, tangents):
+    assert False  # TODO
+
+  def transpose(_, *args):
+    assert False  # TODO
+box_get_p = BoxGet('box_get')
+
+def box_get(box):
+  leaf_vals = box_get_p.bind(box)
+  return jax.tree.unflatten(core.typeof(box)._treedef, leaf_vals)
+
+
+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)
+    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 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 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)
+
+  def test_type_changing_box(self):
+    box = Box(jnp.arange(1))
+    box_set(box, jnp.arange(2))
+    self.assertLen(box._val, 2)
+
+    @jax.jit
+    def f(box, x):
+      box_set(box, x)
+
+    f(box, jnp.arange(3))
+    self.assertLen(box._val, 3)
+    f(box, jnp.arange(4))
+    self.assertLen(box._val, 4)
+
+  def test_pytree_box(self):
+    box = Box(None)
+
+    @jax.jit
+    def f(box, x):
+      assert tracing_ok
+      val = box_get(box)
+      if val is None:
+        box_set(box, x)
+      else:
+        box_set(box, [x, x])
+
+    tracing_ok = True
+    f(box, 1.0)
+    self.assertAllClose(box_get(box), 1.0, check_dtypes=False)
+    f(box, 2.0)
+    self.assertAllClose(box_get(box), [2.0, 2.0], check_dtypes=False)
+    f(box, 3.0)
+    self.assertAllClose(box_get(box), [3.0, 3.0], check_dtypes=False)
+    tracing_ok = False
+    f(box, 4.0)
+    self.assertAllClose(box_get(box), [4.0, 4.0], check_dtypes=False)
+
+  def test_pytree_of_hijaxtypes_box(self):
+
+    @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)
+
+    core.pytype_aval_mappings[MyArray] = lambda _: MyTy()
+
+    box = Box([MyArray(jnp.float32(1)),
+               MyArray(jnp.float32(2))])
+
+    @jax.jit
+    def f(box):
+      a, b = box_get(box)
+      box_set(box, [b, a])
+
+    f(box)
+    val = box_get(box)
+    self.assertIsInstance(val, list)
+    self.assertLen(val, 2)
+    b_, a_ = val
+    self.assertIsInstance(a_, MyArray)
+    self.assertIsInstance(b_, MyArray)
+    self.assertAllClose(a_.arr, 1, check_dtypes=False)
+    self.assertAllClose(b_.arr, 2, check_dtypes=False)
+
+
+if __name__ == '__main__':
+  absltest.main(testLoader=jtu.JaxTestLoader())