# Copyright 2022 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations from collections.abc import Callable, Sequence from functools import partial import itertools as it from typing import Any, NamedTuple from absl.testing import absltest from absl.testing import parameterized import hypothesis as hp import hypothesis.extra.numpy as hnp import hypothesis.strategies as hps import jax from jax import api_util from jax import lax from jax import random from jax._src import config from jax._src import core from jax._src import dtypes from jax._src import hypothesis_test_util as htu from jax._src import linear_util as lu from jax._src import test_util as jtu from jax._src import flattree as ft from jax._src.interpreters import partial_eval as pe from jax._src.state import types as state_types from jax._src.state.discharge import ( discharged_aval, discharge_state, run_state, run_state_reference, ) from jax._src.state.primitives import ( addupdate_p,get_p, pin, ref_addupdate, ref_get, ref_set, ref_swap, swap_p, unpin, with_memory_space_constraint) from jax._src.state.types import ( AbstractRef, AccumEffect, ReadEffect, WriteEffect,shaped_array_ref) from jax._src.util import tuple_insert import jax.numpy as jnp from jax.sharding import AxisType, PartitionSpec as P import numpy as np config.parse_flags_with_absl() htu.setup_hypothesis() def trace_to_jaxpr(f: Callable, *args: Any): return pe.trace_to_jaxpr( f, ft.flatten_args(*args), api_util.debug_info("state_test", f, (0,) * len(args), {}), ) def wrap_init(f: Callable, nr_args: int): # wrapper for lu.wrap_init with debugging info return lu.wrap_init( f, debug_info=api_util.debug_info("state_test", f, (0,) * nr_args, {})) _cond = lambda fn: lambda *args: jax.lax.cond(jnp.array(True), fn, fn, *args) def _shard_map(fn, out_specs): mesh = jax.make_mesh((jax.device_count(),), ("x",), axis_types=(AxisType.Explicit,)) return jax.shard_map(fn, out_specs=out_specs, mesh=mesh) class StatePrimitivesTest(jtu.JaxTestCase): def test_get_abstract_aval_must_take_in_refs(self): ref_aval = core.ShapedArray((), jnp.float32) def f(x_ref): return [ref_get(x_ref, ())] with self.assertRaises(ValueError): trace_to_jaxpr(f, ref_aval) @parameterized.named_parameters( dict(testcase_name="trivial_get", ref_shape=(1, 2), ref_dtype=jnp.float32, idx=(), out_shape=(1, 2), out_dtype=jnp.float32), dict(testcase_name="get_with_index", ref_shape=(1, 2), ref_dtype=jnp.float32, idx=(0,), out_shape=(2,), out_dtype=jnp.float32), dict(testcase_name="get_with_nonleading_index", ref_shape=(1, 2), ref_dtype=jnp.float32, idx=(slice(None), 0), out_shape=(1,), out_dtype=jnp.float32), dict(testcase_name="get_with_array_index", ref_shape=(1, 2, 3, 4), ref_dtype=jnp.float32, idx=(np.array([0, 1]),), out_shape=(2, 2, 3, 4), out_dtype=jnp.float32), dict(testcase_name="get_with_multiple_array_index", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(np.array([0, 1]), np.array([0, 1])), out_shape=(2, 2, 4), out_dtype=jnp.float32), dict(testcase_name="get_with_nonleading_multiple_array_index", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(None), np.array([0, 1]), slice(None), np.array([0, 1])), out_shape=(2, 1, 2), out_dtype=jnp.float32), dict(testcase_name="get_with_nontrivial_slice", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(0, 1), np.array([0, 1]), slice(None), np.array([0, 1])), out_shape=(2, 1, 2), out_dtype=jnp.float32), dict(testcase_name="get_with_nontrivial_slice2", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(0, 1), slice(1, 3), slice(None), slice(None)), out_shape=(1, 2, 2, 4), out_dtype=jnp.float32), dict(testcase_name="get_with_ref_simple_at", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(1, 3), slice(None), slice(None)), out_shape=(2, 2, 4), out_dtype=jnp.float32, at_indices=((0,),)), dict(testcase_name="get_with_ref_simple_at2", ref_shape=(6, 1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(0, 2), slice(0, 1), slice(1, 3), slice(None), slice(None)), out_shape=(2, 1, 2, 2, 4), out_dtype=jnp.float32, at_indices=((slice(2, 6),),)), dict(testcase_name="get_with_ref_multiple_at", ref_shape=(1, 3, 5, 4), ref_dtype=jnp.float32, idx=(slice(None), slice(None), slice(0, 2)), out_shape=(3, 1, 2), out_dtype=jnp.float32, at_indices=((0,), (slice(None), slice(0, 1)))), ) def test_get_abstract_eval(self, ref_shape, ref_dtype, idx, out_shape=None, out_dtype=None, at_indices=(), should_error=False): ref_aval = AbstractRef(core.ShapedArray(ref_shape, ref_dtype)) def f(x_ref): for at_idx in at_indices: x_ref = x_ref.at[at_idx] out = ref_get(x_ref, idx) return [out] if should_error: with self.assertRaises(Exception): trace_to_jaxpr(f, ref_aval) else: jaxpr, out_avals_ft = trace_to_jaxpr(f, ref_aval) out_avals = list(out_avals_ft.unflatten()) self.assertSetEqual( jaxpr.effects, {ReadEffect(jaxpr.invars[0])} ) self.assertLen(out_avals, 1) out_aval, = out_avals self.assertIsInstance(out_aval, core.ShapedArray) self.assertEqual(out_aval.shape, out_shape) self.assertEqual(out_aval.dtype, out_dtype) @parameterized.parameters( ((4, 5), 0, (0,)), ((4, 5), 1, (0,)), ((9, 10, 11, 12), 0, (slice(None), 0, 1)), # Contiguous int indexing ((9, 10, 11, 12), 0, (0, slice(None), 1)), # Non-contiguous int indexing ((9, 10, 11, 12), 1, (slice(None), 0, 1)), # Contiguous after batch ((9, 10, 11, 12), 2, (slice(None), 0, 1)), # Non-contiguous after batch ((9, 10, 11, 12), 3, (slice(None), slice(None), 0)), # Shaped int indexer, contiguous after batch ((9, 10, 11, 12), 3, (slice(None), slice(None), np.array([[0,1]]))), # Shaped int indexer, non-contiguous after batch ((9, 10, 11, 12), 2, (np.array([[0, 1]]), slice(None), np.array([[0, 1]]))), ) def test_vmap_of_get_regression(self, shape, in_axes, indexer): # Regression test for https://github.com/jax-ml/jax/issues/33309 def f(x): return x[indexer] x = jnp.ones(shape) result = jax.vmap(f, in_axes=in_axes)(jax.new_ref(x)) expected = jax.vmap(f, in_axes=in_axes)(x) self.assertArraysEqual(result, expected) def test_swap_abstract_eval_must_take_in_refs(self): ref_aval = core.ShapedArray((), jnp.float32) val_aval = core.ShapedArray((), jnp.float32) def f(x_ref, val): return [ref_swap(x_ref, (), val)] with self.assertRaises(ValueError): trace_to_jaxpr(f, ref_aval, val_aval) @parameterized.named_parameters( dict(testcase_name="invalid_val_shape", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(), should_error=True), dict(testcase_name="invalid_val_shape_slice", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(slice(None),), should_error=True), dict(testcase_name="trivial_swap", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(1, 2), val_dtype=jnp.float32, idx=(), out_shape=(1, 2), out_dtype=jnp.float32), dict(testcase_name="bad_dtype", ref_shape=(1, 2), ref_dtype=jnp.int32, val_shape=(1, 2), val_dtype=jnp.float32, idx=(), should_error=True), dict(testcase_name="swap_with_index", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(0,), out_shape=(2,), out_dtype=jnp.float32), dict(testcase_name="swap_with_nonleading_index", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(1,), val_dtype=jnp.float32, idx=(slice(None), 0), out_shape=(1,), out_dtype=jnp.float32), dict(testcase_name="swap_with_nonleading_index_bad_val", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(slice(None), 0), should_error=True), dict(testcase_name="swap_with_array_index", ref_shape=(1, 2, 3, 4), ref_dtype=jnp.float32, val_shape=(2, 2, 3, 4), val_dtype=jnp.float32, idx=(np.array([0, 1]),), out_shape=(2, 2, 3, 4), out_dtype=jnp.float32), dict(testcase_name="swap_with_multiple_array_index", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 2, 4), val_dtype=jnp.float32, idx=(np.array([0, 1]), np.array([0, 1])), out_shape=(2, 2, 4), out_dtype=jnp.float32), dict(testcase_name="swap_with_nonleading_multiple_array_index", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 1, 2), val_dtype=jnp.float32, idx=(slice(None), np.array([0, 1]), slice(None), np.array([0, 1])), out_shape=(2, 1, 2), out_dtype=jnp.float32), dict(testcase_name="swap_with_nontrivial_slice", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(0, 1), np.array([0, 1]), slice(None), np.array([0, 1])), val_shape=(2, 1, 2), val_dtype=jnp.float32, out_shape=(2, 1, 2), out_dtype=jnp.float32), dict(testcase_name="swap_with_nontrivial_slice2", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(0, 1), slice(1, 3), slice(None), slice(None)), val_shape=(1, 2, 2, 4), val_dtype=jnp.float32, out_shape=(1, 2, 2, 4), out_dtype=jnp.float32), dict(testcase_name="swap_with_ref_simple_at", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(0, 1), slice(1, 3), slice(None),), val_shape=(1, 1, 4), val_dtype=jnp.float32, out_shape=(1, 1, 4), out_dtype=jnp.float32, at_indices=((0,),),), dict(testcase_name="swap_with_ref_simple_at2", ref_shape=(4, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(None), slice(0, 1), slice(1, 3), slice(None),), val_shape=(2, 1, 1, 4), val_dtype=jnp.float32, out_shape=(2, 1, 1, 4), out_dtype=jnp.float32, at_indices=((slice(0, 2),),),), dict(testcase_name="swap_with_ref_multiple_at2", ref_shape=(1, 4, 3, 2, 4), ref_dtype=jnp.float32, idx=(slice(None), slice(0, 1), slice(1, 3), slice(None),), val_shape=(2, 1, 1, 4), val_dtype=jnp.float32, out_shape=(2, 1, 1, 4), out_dtype=jnp.float32, at_indices=((slice(None), slice(0, 2),), (0,)),), ) def test_swap_abstract_eval(self, ref_shape, ref_dtype, val_shape, val_dtype, idx, out_shape=None, out_dtype=None, at_indices=(), should_error=False): ref_aval = AbstractRef(core.ShapedArray(ref_shape, ref_dtype)) val_aval = core.ShapedArray(val_shape, val_dtype) def f(x_ref, val): for at_idx in at_indices: x_ref = x_ref.at[at_idx] out = ref_swap(x_ref, idx, val) return [out] if should_error: with self.assertRaises(Exception): trace_to_jaxpr(f, ref_aval, val_aval) else: jaxpr, out_avals_ft = trace_to_jaxpr(f, ref_aval, val_aval) out_avals = list(out_avals_ft.unflatten()) self.assertSetEqual( jaxpr.effects, {WriteEffect(jaxpr.invars[0])} ) self.assertLen(out_avals, 1) out_aval, = out_avals self.assertIsInstance(out_aval, core.ShapedArray) self.assertEqual(out_aval.shape, out_shape) self.assertEqual(out_aval.dtype, out_dtype) @parameterized.named_parameters( dict(testcase_name="invalid_val_shape", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(), should_error=True), dict(testcase_name="invalid_val_shape_slice", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(slice(None),), should_error=True), dict(testcase_name="trivial_addupdate", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(1, 2), val_dtype=jnp.float32, idx=(),), dict(testcase_name="bad_dtype", ref_shape=(1, 2), ref_dtype=jnp.int32, val_shape=(1, 2), val_dtype=jnp.float32, idx=(), should_error=True), dict(testcase_name="addupdate_with_index", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(0,),), dict(testcase_name="addupdate_with_nonleading_index", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(1,), val_dtype=jnp.float32, idx=(slice(None), 0)), dict(testcase_name="addupdate_with_nonleading_index_bad_val", ref_shape=(1, 2), ref_dtype=jnp.float32, val_shape=(2,), val_dtype=jnp.float32, idx=(slice(None), 0), should_error=True), dict(testcase_name="addupdate_with_array_index", ref_shape=(1, 2, 3, 4), ref_dtype=jnp.float32, val_shape=(2, 2, 3, 4), val_dtype=jnp.float32, idx=(np.array([0, 1]),)), dict(testcase_name="addupdate_with_multiple_array_index", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 2, 4), val_dtype=jnp.float32, idx=(np.array([0, 1]), np.array([0, 1]))), dict(testcase_name="addupdate_with_nonleading_multiple_array_index", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 1, 2), val_dtype=jnp.float32, idx=(slice(None), np.array([0, 1]), slice(None), np.array([0, 1]))), dict(testcase_name="ref_with_simple_at", ref_shape=(1, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 2), val_dtype=jnp.float32, idx=(np.array([0, 1]), slice(None), np.array([0, 1])), at_indices=((0,),)), dict(testcase_name="ref_with_simple_at2", ref_shape=(3, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 3, 4), val_dtype=jnp.float32, idx=(np.array([0, 1]), slice(None), np.array([0, 1])), at_indices=((slice(0, 3),),)), dict(testcase_name="ref_with_multiple_at", ref_shape=(3, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 2), val_dtype=jnp.float32, idx=(np.array([0, 1]), slice(None), np.array([0, 1])), at_indices=((slice(0, 3),), (0,))), dict(testcase_name="ref_with_multiple_at2", ref_shape=(3, 3, 2, 4), ref_dtype=jnp.float32, val_shape=(2, 2), val_dtype=jnp.float32, idx=(np.array([0, 1]), slice(None), np.array([0, 1])), at_indices=((slice(None), slice(0, 3),), (0,))), ) def test_addupdate_abstract_eval(self, ref_shape, ref_dtype, val_shape, val_dtype, idx, at_indices=(), should_error=False): ref_aval = AbstractRef(core.ShapedArray(ref_shape, ref_dtype)) val_aval = core.ShapedArray(val_shape, val_dtype) def f(x_ref, val): for at_idx in at_indices: x_ref = x_ref.at[at_idx] ref_addupdate(x_ref, idx, val) return [] if should_error: with self.assertRaises(Exception): trace_to_jaxpr(f, ref_aval, val_aval) else: jaxpr, out_avals_ft = trace_to_jaxpr(f, ref_aval, val_aval) out_avals = list(out_avals_ft.unflatten()) self.assertSetEqual( jaxpr.effects, {AccumEffect(jaxpr.jaxpr.invars[0])} ) self.assertLen(out_avals, 0) def test_addupdate_abstract_eval_must_take_in_refs(self): ref_aval = core.ShapedArray((), jnp.float32) val_aval = core.ShapedArray((), jnp.float32) def f(x_ref, val): return [ref_addupdate(x_ref, (), val)] with self.assertRaises(ValueError): trace_to_jaxpr(f, ref_aval, val_aval) def test_can_represent_get_and_swap_in_jaxprs(self): def body(x): x[()] = jnp.int32(1) x[()] = jnp.int32(2) return (x[()],) jaxpr, out_avals_ft = trace_to_jaxpr(body, shaped_array_ref((), jnp.int32)) out_avals = list(out_avals_ft.unflatten()) self.assertLen(jaxpr.consts, 0) self.assertListEqual(out_avals, [core.ShapedArray((), jnp.int32)]) self.assertEqual(jaxpr.eqns[0].primitive, swap_p) self.assertEqual(jaxpr.eqns[1].primitive, swap_p) self.assertEqual(jaxpr.eqns[2].primitive, get_p) def test_can_represent_addupdate_in_jaxprs(self): def body(x): ref_addupdate(x, (), jnp.int32(1)) return (x[()],) jaxpr, out_avals_ft = trace_to_jaxpr(body, shaped_array_ref((), jnp.int32)) out_avals = list(out_avals_ft.unflatten()) self.assertLen(jaxpr.consts, 0) self.assertListEqual(out_avals, [core.ShapedArray((), jnp.int32)]) self.assertEqual(jaxpr.eqns[0].primitive, addupdate_p) def test_get_custom_pretty_printing_rule(self): def body(x_ref): x = x_ref[()] return [x] jaxpr, _ = trace_to_jaxpr(body, shaped_array_ref((), jnp.int32)) self.assertIn( "b:i32[] <- a[...]", jaxpr.pretty_print(use_color=False) ) def body(x_ref): x = x_ref[:, 0] return [x] jaxpr, _ = trace_to_jaxpr(body, shaped_array_ref((1, 2), jnp.int32)) self.assertIn( "b:i32[1] <- a[:,0]", jaxpr.pretty_print(use_color=False) ) def test_set_custom_pretty_printing_rule(self): def body(x_ref): x_ref[()] = jnp.int32(2) return [] jaxpr, _ = trace_to_jaxpr(body, shaped_array_ref((), jnp.int32)) self.assertIn( "a[...] <- 2:i32[]", jaxpr.pretty_print(use_color=False) ) def body(x_ref, val): x_ref[:, 0] = val return [] jaxpr, _ = trace_to_jaxpr( body, shaped_array_ref((1, 2), jnp.int32), core.ShapedArray((1,), jnp.int32), ) self.assertIn("a[:,0] <- b", jaxpr.pretty_print(use_color=False)) def test_swap_custom_pretty_printing_rule(self): def body(x_ref): x = ref_swap(x_ref, (), jnp.int32(2)) return [x] jaxpr, _ = trace_to_jaxpr(body, shaped_array_ref((), jnp.int32)) self.assertIn( "b:i32[], a[...] <- a[...], 2:i32[]", jaxpr.pretty_print(use_color=False), ) def body(x_ref, val): x = ref_swap(x_ref, (slice(None), 0), val) return [x] jaxpr, _ = trace_to_jaxpr( body, shaped_array_ref((1, 2), jnp.int32), core.ShapedArray((1,), jnp.int32), ) self.assertIn( "c:i32[1], a[:,0] <- a[:,0], b", jaxpr.pretty_print(use_color=False), ) def test_addupdate_custom_pretty_printing_rule(self): def body(x_ref): ref_addupdate(x_ref, (), jnp.int32(2)) return [] jaxpr, _ = trace_to_jaxpr(body, shaped_array_ref((), jnp.int32)) self.assertIn("a[...] += 2", jaxpr.pretty_print(use_color=False)) def body(x_ref, val): ref_addupdate(x_ref, (slice(None), 0), val) return [] jaxpr, _ = trace_to_jaxpr( body, shaped_array_ref((1, 2), jnp.int32), core.ShapedArray((1,), jnp.int32), ) self.assertIn("a[:,0] += b", jaxpr.pretty_print(use_color=False)) def test_get_jvp(self): def f(r): x = r[()] return jnp.cos(x) def g(r, rdot): return jax.jvp(f, (r,), (rdot,)) in_avals = [shaped_array_ref((), jnp.dtype('float32')), shaped_array_ref((), jnp.dtype('float32'))] jaxpr, _ = trace_to_jaxpr(g, *in_avals) self.assertEqual(jaxpr.eqns[0].primitive, get_p) self.assertEqual(jaxpr.eqns[1].primitive, get_p) def test_swap_jvp(self): def f(a): x = a[()] a[()] = jnp.sin(x) return a[()] def g(r, rdot): return jax.jvp(f, (r,), (rdot,)) in_avals = [shaped_array_ref((), jnp.dtype('float32')), shaped_array_ref((), jnp.dtype('float32'))] jaxpr, _ = trace_to_jaxpr(g, *in_avals) self.assertEqual(jaxpr.eqns[0].primitive, get_p) self.assertEqual(jaxpr.eqns[1].primitive, get_p) self.assertEqual(jaxpr.eqns[2].primitive, lax.sin_p) self.assertEqual(jaxpr.eqns[3].primitive, lax.cos_p) self.assertEqual(jaxpr.eqns[4].primitive, lax.mul_p) self.assertEqual(jaxpr.eqns[5].primitive, swap_p) self.assertEqual(jaxpr.eqns[6].primitive, swap_p) def test_addupdate_jvp(self): def f(a): ref_addupdate(a, (), jnp.float32(1.)) return a[()] def g(r, rdot): return jax.jvp(f, (r,), (rdot,)) in_avals = [shaped_array_ref((), jnp.dtype('float32')), shaped_array_ref((), jnp.dtype('float32'))] jaxpr, _ = trace_to_jaxpr(g, *in_avals) self.assertEqual(jaxpr.eqns[0].primitive, addupdate_p) self.assertEqual(jaxpr.eqns[1].primitive, addupdate_p) self.assertEqual(jaxpr.eqns[2].primitive, get_p) self.assertEqual(jaxpr.eqns[3].primitive, get_p) @jtu.sample_product( [dict(ref_shape=ref_shape, ref_bdim=ref_bdim, idx_shape=idx_shape, indexed_dims=indexed_dims, idx_bdims=idx_bdims, out_bdim=out_bdim) for ref_shape in [(1,), (2, 3), (4, 5, 6)] for ref_bdim in range(1 + len(ref_shape)) for idx_shape in [(), (1,), (2,), (5, 6)] for indexed_dims in it.product([True, False], repeat=len(ref_shape)) for idx_bdims in it.product([None, *range(1 + len(idx_shape))], repeat=sum(indexed_dims)) for out_bdim in range(1 + len(ref_shape) - sum(indexed_dims) + len(idx_shape) * any(indexed_dims)) ], op=[ lambda x_ref, indexer: [x_ref[indexer]], lambda x_ref, indexer: [ ref_swap(x_ref, indexer, jnp.ones_like(x_ref[indexer]))], lambda x_ref, indexer: ( ref_addupdate(x_ref, indexer, jnp.ones_like(x_ref[indexer])) or [jnp.ones_like(x_ref[indexer])]), ], ) def test_vmap(self, ref_shape, ref_bdim, idx_shape, indexed_dims, idx_bdims, out_bdim, op): intx = dtypes.default_int_dtype() floatx = dtypes.default_float_dtype() axis_size = 7 def maybe_insert(shape, idx): if idx is None: return shape return tuple_insert(shape, idx, axis_size) batched_ref_shape = maybe_insert(ref_shape, ref_bdim) ref_aval = shaped_array_ref(ref_shape, floatx) bat_ref_aval = shaped_array_ref(batched_ref_shape, floatx) idx_avals = [core.ShapedArray(idx_shape, intx) for _ in idx_bdims] bat_idx_avals = [ core.ShapedArray(maybe_insert(idx_shape, idx_bdim), intx) for idx_bdim in idx_bdims] def f(x_ref, *idxs): idxs_ = iter(idxs) indexer = tuple(next(idxs_) if b else slice(None) for b in indexed_dims) return op(x_ref, indexer) rng = self.rng() a = rng.randn(*bat_ref_aval.shape).astype(floatx) his = [d for d, b in zip(ref_aval.shape, indexed_dims) if b] idxs = [rng.randint(low=0, high=hi, size=i.shape, dtype=intx) for i, hi in zip(bat_idx_avals, his)] # discharge-of-vmap f_batched = jax.vmap(f, in_axes=(ref_bdim, *idx_bdims), out_axes=[out_bdim]) jaxpr, _ = trace_to_jaxpr(f_batched, bat_ref_aval, *bat_idx_avals) jaxpr = discharge_state(jaxpr) discharge_of_vmap_ans = core.eval_jaxpr(jaxpr, jaxpr.consts, a, *idxs) # vmap-of-discharge jaxpr, _ = trace_to_jaxpr(f, ref_aval, *idx_avals) closed_jaxpr_ = discharge_state(jaxpr) f_batched = jax.vmap(partial(core.eval_jaxpr, closed_jaxpr_.jaxpr, closed_jaxpr_.consts), in_axes=(ref_bdim, *idx_bdims), out_axes=[out_bdim, ref_bdim]) vmap_of_discharge_ans = f_batched(a, *idxs) self.assertAllClose(discharge_of_vmap_ans, vmap_of_discharge_ans, check_dtypes=False) class StateDischargeTest(jtu.JaxTestCase): def test_discharge_get(self): def f(a_ref): a = ref_get(a_ref, ()) return [a + 1] in_avals = [shaped_array_ref((), jnp.dtype('float32'))] jaxpr, _ = trace_to_jaxpr(f, *in_avals) # Discharging should just turn this into a jaxpr that just adds 1. discharged_jaxpr = discharge_state(jaxpr) self.assertLen(discharged_jaxpr.invars, 1) self.assertLen(discharged_jaxpr.outvars, 2) self.assertEqual(discharged_jaxpr.eqns[0].primitive, lax.add_p) # Should be able to evaluate this jaxpr self.assertListEqual(core.eval_jaxpr(discharged_jaxpr, discharged_jaxpr.consts, jnp.float32(1.)), [2., 1.]) def test_discharge_get_with_slice(self): def f(a_ref): a = ref_get(a_ref, (0, 1)) return [a + 1] in_avals = [shaped_array_ref((4, 3, 2), jnp.dtype('float32'))] jaxpr, _ = trace_to_jaxpr(f, *in_avals) # Discharging should just turn this into a jaxpr that just adds 1. discharged_jaxpr = discharge_state(jaxpr) self.assertLen(discharged_jaxpr.invars, 1) self.assertLen(discharged_jaxpr.outvars, 2) self.assertIn(lax.dynamic_slice_p, {eqn.primitive for eqn in discharged_jaxpr.eqns}) # Should be able to evaluate this jaxpr inval = jnp.arange(24., dtype=jnp.float32).reshape((4, 3, 2)) outval, refval = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((outval == inval[0, 1] + 1).all()) self.assertTrue((refval == inval).all()) def test_discharge_get_with_gather(self): def f(a_ref): a = a_ref[jnp.array([0, 1])] return [a + 1] in_avals = [shaped_array_ref((4, 3), jnp.dtype('float32'))] jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(jaxpr) inval = jnp.arange(4 * 3, dtype=jnp.float32).reshape((4, 3)) outval, refval = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((outval == inval[jnp.array([0, 1])] + 1).all()) self.assertTrue((refval == inval).all()) def test_discharge_set(self): def f(a_ref, b): ref_set(a_ref, (), b + 1) return [] in_avals = [shaped_array_ref((), jnp.dtype('float32')), core.ShapedArray((), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) # Discharging should just turn this into a jaxpr that ignores the first # value and returns second value plus 1. discharged_jaxpr = discharge_state(stateful_jaxpr) self.assertLen(discharged_jaxpr.invars, 2) self.assertLen(discharged_jaxpr.outvars, 1) self.assertEqual(core.eval_jaxpr(discharged_jaxpr, discharged_jaxpr.consts, jnp.float32(0.), jnp.float32(1.))[0], 2.) self.assertEqual(core.eval_jaxpr(discharged_jaxpr, discharged_jaxpr.consts, jnp.float32(2.), jnp.float32(1.))[0], 2.) def test_discharge_set_with_slice(self): def f(a_ref): ref_set(a_ref, (0, 1), jnp.ones(2, dtype=jnp.dtype('float32'))) return [] in_avals = [shaped_array_ref((4, 3, 2), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) # Discharging should just turn this into a jaxpr that just adds 1. discharged_jaxpr = discharge_state(stateful_jaxpr) self.assertLen(discharged_jaxpr.invars, 1) self.assertLen(discharged_jaxpr.outvars, 1) self.assertIn(lax.dynamic_update_slice_p, {eqn.primitive for eqn in discharged_jaxpr.eqns}) self.assertIn(lax.dynamic_slice_p, {eqn.primitive for eqn in discharged_jaxpr.eqns}) # Should be able to evaluate this jaxpr inval = jnp.arange(24., dtype=jnp.float32).reshape((4, 3, 2)) refval, = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((refval == inval.at[0, 1].set(1.)).all()) def test_discharge_set_with_gather(self): def f(a_ref): a_ref[jnp.array([0, 1])] = jnp.ones((2, 3), 'float32') return [] in_avals = [shaped_array_ref((4, 3), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) inval = jnp.arange(4 * 3, dtype=jnp.float32).reshape((4, 3)) refval, = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((refval == inval.at[jnp.array([0, 1])].set(1.)).all()) def test_discharge_swap(self): def f(a_ref): a = ref_swap( a_ref.at[0:4, 0:3, 0:2].at[1:3, :, 0], (slice(None), slice(1, 3)), jnp.zeros((2, 2), jnp.float32)) return [a + 1] in_avals = [shaped_array_ref((4, 3, 2), jnp.float32)] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) self.assertLen(discharged_jaxpr.invars, 1) self.assertLen(discharged_jaxpr.outvars, 2) inval = jnp.arange(24., dtype=jnp.float32).reshape((4, 3, 2)) outval, refval = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertArraysEqual(outval, inval[1:3, 1:3, 0] + 1) self.assertArraysEqual(refval, inval.at[1:3, 1:3, 0].set(0)) def test_discharge_addupdate(self): def f(a_ref, b): ref_addupdate(a_ref, (), b + 1) return [] in_avals = [shaped_array_ref((), jnp.dtype('float32')), core.ShapedArray((), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) # Discharging should just turn this into a jaxpr that adds the first value, # second value, and 1. discharged_jaxpr = discharge_state(stateful_jaxpr) self.assertLen(discharged_jaxpr.invars, 2) self.assertLen(discharged_jaxpr.outvars, 1) self.assertEqual(core.eval_jaxpr(discharged_jaxpr, discharged_jaxpr.consts, jnp.float32(0.), jnp.float32(1.))[0], 2.) self.assertEqual(core.eval_jaxpr(discharged_jaxpr, discharged_jaxpr.consts, jnp.float32(2.), jnp.float32(1.))[0], 4.) def test_discharge_addupdate_with_slice(self): def f(a_ref): ref_addupdate(a_ref, (0, 1), jnp.ones(2, dtype=jnp.dtype('float32'))) return [] in_avals = [shaped_array_ref((4, 3, 2), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) self.assertLen(discharged_jaxpr.invars, 1) self.assertLen(discharged_jaxpr.outvars, 1) self.assertIn(lax.dynamic_update_slice_p, {eqn.primitive for eqn in discharged_jaxpr.eqns}) self.assertIn(lax.add_p, {eqn.primitive for eqn in discharged_jaxpr.eqns}) self.assertIn(lax.dynamic_slice_p, {eqn.primitive for eqn in discharged_jaxpr.eqns}) inval = jnp.arange(24., dtype=jnp.float32).reshape((4, 3, 2)) refval, = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((refval == inval.at[0, 1].add(1.)).all()) def test_discharge_addupdate_with_gather(self): def f(a_ref): ref_addupdate(a_ref, (jnp.array([0, 1]),), jnp.ones((2, 3), 'float32')) return [] in_avals = [shaped_array_ref((4, 3), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) inval = jnp.arange(4 * 3, dtype=jnp.float32).reshape((4, 3)) refval, = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((refval == inval.at[jnp.array([0, 1])].add(1.)).all()) def test_discharge_jaxpr_with_multiple_outputs(self): def f(a_ref): a = ref_get(a_ref, ()) b = a + 1 return [a, b] in_avals = [shaped_array_ref((4,), jnp.dtype('float32'))] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) self.assertLen(discharged_jaxpr.invars, 1) self.assertLen(discharged_jaxpr.outvars, 3) inval = jnp.arange(4., dtype=jnp.float32) a, b, refval = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, inval) self.assertTrue((a == inval).all()) self.assertTrue((b == inval + 1).all()) self.assertTrue((refval == inval).all()) def test_partially_discharging_jaxpr_keeps_refs(self): def f(a_ref, b_ref): ref_set(a_ref, (), jnp.ones(4, jnp.float32)) ref_set(b_ref, (), jnp.ones(4, jnp.float32)) return [] in_avals = [ shaped_array_ref((4,), jnp.dtype('float32')), shaped_array_ref((4,), jnp.dtype('float32')) ] stateful_jaxpr, _ = trace_to_jaxpr(f, *in_avals) discharged_jaxpr = discharge_state(stateful_jaxpr, should_discharge=[False, True]) self.assertLen(discharged_jaxpr.invars, 2) self.assertLen(discharged_jaxpr.outvars, 1) self.assertIsInstance(discharged_jaxpr.invars[0].aval, AbstractRef) self.assertIsInstance(discharged_jaxpr.invars[1].aval, core.ShapedArray) self.assertEqual(discharged_jaxpr.effects, {WriteEffect(discharged_jaxpr.invars[0])}) def test_ellipsis_index(self): def f(ref): ref_set(ref, ..., jnp.array(0., dtype=jnp.float32)) ref_get(ref, ...) ref[...] = jnp.array(0., dtype=jnp.float32) ref[...] return [] in_avals = [shaped_array_ref((), jnp.float32)] trace_to_jaxpr(f, *in_avals) def test_partial_discharge(self): def f(a_ref, b_ref): a_ref[...] = jnp.array(0., dtype=jnp.float32) b_ref[...] = jnp.array(1., dtype=jnp.float32) return a_ref[...], b_ref[...] scalar_ref_1 = shaped_array_ref((), jnp.float32) scalar_ref_2 = shaped_array_ref((), jnp.float32) jaxpr, _ = trace_to_jaxpr(f, scalar_ref_1, scalar_ref_2) discharged_jaxpr = discharge_state( jaxpr, should_discharge=[False, True] ) prim_count = lambda p, j: sum(eqn.primitive == p for eqn in j.eqns) self.assertEqual( prim_count(swap_p, jaxpr) // 2, prim_count(swap_p, discharged_jaxpr), ) self.assertEqual( prim_count(get_p, jaxpr) // 2, prim_count(get_p, discharged_jaxpr) ) def test_partial_fori_discharge(self): def f(a_ref, b_ref): def body(i, st): a_ref[...] += 2 * i b_ref[...] += i return () lax.fori_loop(0, 5, body, init_val=()) return a_ref[...], b_ref[...] ref = lambda x: AbstractRef(core.typeof(x)) f_jaxpr = jax.make_jaxpr(f)(ref(1.), ref(2.)) jaxpr = discharge_state(f_jaxpr, should_discharge=[False, True]) # Effects on y_ref were discharged away but not the effects on x_ref a, b = f_jaxpr.invars self.assertEqual(f_jaxpr.effects, {ReadEffect(a), WriteEffect(a), ReadEffect(b), WriteEffect(b)}) self.assertEqual(jaxpr.effects, {ReadEffect(jaxpr.invars[0]), WriteEffect(jaxpr.invars[0])}) # x_ref arg is still a reference but y_ref is discharged self.assertNotIsInstance(jaxpr.invars[1].aval, AbstractRef) self.assertIsInstance(jaxpr.invars[0].aval, AbstractRef) # x_ref value is returned as part of the discharged refs set. self.assertLen(f_jaxpr.out_avals, 2) self.assertLen(jaxpr.outvars, 3) def test_discharge_call_primitive(self): def f(y_ref, x): def g(x): y_ref[...] = x + 1.0 return [] jaxpr, _, consts = pe.trace_to_jaxpr_dynamic( wrap_init(g, 1), [core.typeof(x)]) return core.call_p.bind(*consts, x, call_jaxpr=jaxpr) y_ref_aval = shaped_array_ref((), jnp.float32) x_aval = core.ShapedArray((), jnp.float32) stateful_jaxpr, _ = trace_to_jaxpr(f, y_ref_aval, x_aval) # This should successfully discharge without a TypeError. discharge_state(stateful_jaxpr, should_discharge=[True, False]) def index_arrays(size, idx_shape): valid_idx = hps.integers(min_value=-size, max_value=size - 1) return hnp.arrays(np.int32, idx_shape, elements=valid_idx) Shape = tuple[int, ...] class IndexParam(NamedTuple): ref_aval: shaped_array_ref ref_shape: Shape indexed_dims: list[bool] idx_avals: tuple[core.ShapedArray, ...] idx_shape: Shape slice_aval: core.ShapedArray slice_shape: Shape @hps.composite def index_params(draw): ref_shape = draw(hnp.array_shapes(max_dims=4, max_side=7), label='ref_shape') indexed_dims = draw(hps.lists(hps.booleans(), min_size=len(ref_shape), max_size=len(ref_shape))) idx_shape = draw(hnp.array_shapes(min_dims=0, max_dims=3, max_side=5)) if not any(indexed_dims): slice_shape = ref_shape else: sliced_shape = tuple(s for s, b in zip(ref_shape, indexed_dims) if not b) int_indexers_contiguous = bool( np.all(np.diff(np.where(indexed_dims)[0]) == 1) ) if not int_indexers_contiguous: slice_shape = (*idx_shape, *sliced_shape) else: insert_pos = indexed_dims.index(True) slice_shape = ( *sliced_shape[:insert_pos], *idx_shape, *sliced_shape[insert_pos:], ) ref_aval = shaped_array_ref(ref_shape, np.float32) idx_avals = tuple(core.ShapedArray(idx_shape, np.int32) for _ in range(sum(indexed_dims))) slice_aval = core.ShapedArray(slice_shape, np.float32) return IndexParam(ref_aval, ref_shape, indexed_dims, idx_avals, idx_shape, slice_aval, slice_shape) class VmappableIndexParam(NamedTuple): index_param: IndexParam ref_bdim: int | None non_slice_idx_bdims: tuple[int | None, ...] slice_bdim: int bat_ref_aval: shaped_array_ref bat_ref_shape: Shape bat_non_slice_idx_avals: tuple[core.ShapedArray, ...] bat_non_slice_idx_shapes: tuple[Shape, ...] bat_slice_aval: core.ShapedArray bat_slice_shape: Shape def maybe_tuple_insert(t: tuple[Any, ...], idx: int | None, val: Any) -> tuple[Any, ...]: if idx is None: return t return tuple_insert(t, idx, val) @hps.composite def vmappable_index_params(draw, *, op_type: str): axis_size = draw(hps.integers(min_value=1, max_value=7), label='axis_size') index_param: IndexParam = draw(index_params()) non_slice_idx_bdims = tuple( draw(hps.one_of( hps.none(), hps.integers(min_value=0, max_value=len(index_param.idx_shape)))) for b in index_param.indexed_dims if b) bat_non_slice_idx_shapes = tuple( maybe_tuple_insert(index_param.idx_shape, idx_bdim, axis_size) for idx_bdim in non_slice_idx_bdims) if op_type == "swap": # In a swap, the ref *must* be batched ref_bdim = draw(hps.integers(min_value=0, max_value=len(index_param.ref_shape))) if any(idx_bdim is not None for idx_bdim in non_slice_idx_bdims): # If it's a swap, if indices are batched, val must be batched. slice_bdim = draw(hps.integers( min_value=0, max_value=len(index_param.slice_shape))) else: slice_bdim = draw(hps.one_of(hps.none(), hps.integers( min_value=0, max_value=len(index_param.slice_shape)))) elif op_type == "get": # In a get, the indices must be batched or ref is batched if all(idx_bdim is None for idx_bdim in non_slice_idx_bdims): ref_bdim = draw(hps.integers(min_value=0, max_value=len(index_param.ref_shape))) else: ref_bdim = draw(hps.one_of(hps.none(), hps.integers(min_value=0, max_value=len(index_param.ref_shape)))) slice_bdim = draw(hps.integers( min_value=0, max_value=len(index_param.slice_shape))) bat_ref_shape = maybe_tuple_insert(index_param.ref_shape, ref_bdim, axis_size) bat_ref_aval = shaped_array_ref(bat_ref_shape, np.float32) bat_non_slice_idx_avals = tuple( core.ShapedArray(shape, np.int32) for shape in bat_non_slice_idx_shapes) bat_slice_shape = maybe_tuple_insert(index_param.slice_shape, slice_bdim, axis_size) bat_slice_aval = core.ShapedArray(bat_slice_shape, np.float32) return VmappableIndexParam(index_param, ref_bdim, non_slice_idx_bdims, slice_bdim, bat_ref_aval, bat_ref_shape, bat_non_slice_idx_avals, bat_non_slice_idx_shapes, bat_slice_aval, bat_slice_shape) class GetVmapParams(NamedTuple): vmap_index_param: VmappableIndexParam bat_ref: np.ndarray bat_idxs: tuple[np.ndarray, ...] def __repr__(self): return ( f"GetVmapParams(vmap_index_param={self.vmap_index_param}, " f"bat_ref=array(shape={self.bat_ref.shape}, dtype={self.bat_ref.dtype}), " f"bat_idxs={tuple(f'array(shape={x.shape}, dtype={x.dtype})' for x in self.bat_idxs)})" ) _float_elements = hps.floats(min_value=-100.0, max_value=100.0, allow_nan=False, allow_infinity=False) @hps.composite def get_vmap_params(draw): vmap_index_param: VmappableIndexParam = draw( vmappable_index_params(op_type="get")) bat_ref = draw(hnp.arrays(np.float32, vmap_index_param.bat_ref_shape, elements=_float_elements)) bat_idx_shapes_ = iter(vmap_index_param.bat_non_slice_idx_shapes) bat_idxs = tuple( draw(index_arrays(size, next(bat_idx_shapes_))) for size, indexed in zip( vmap_index_param.index_param.ref_shape, vmap_index_param.index_param.indexed_dims) if indexed) assert next(bat_idx_shapes_, None) is None return GetVmapParams(vmap_index_param, bat_ref, bat_idxs) class SetVmapParams(NamedTuple): vmap_index_param: VmappableIndexParam bat_ref: np.ndarray bat_val: np.ndarray bat_idxs: tuple[np.ndarray, ...] def __repr__(self): return ( f"SetVmapParams(vmap_index_param={self.vmap_index_param}, " f"bat_ref=array(shape={self.bat_ref.shape}, dtype={self.bat_ref.dtype}), " f"bat_val=array(shape={self.bat_val.shape}, dtype={self.bat_val.dtype}), " f"bat_idxs={tuple(f'array(shape={x.shape}, dtype={x.dtype})' for x in self.bat_idxs)})" ) @hps.composite def set_vmap_params(draw): vmap_index_param: VmappableIndexParam = draw(vmappable_index_params( op_type="swap")) bat_ref = draw(hnp.arrays(np.float32, vmap_index_param.bat_ref_shape, elements=_float_elements)) bat_idx_shapes_ = iter(vmap_index_param.bat_non_slice_idx_shapes) bat_idxs = tuple( draw(index_arrays(size, next(bat_idx_shapes_))) for size, indexed in zip( vmap_index_param.index_param.ref_shape, vmap_index_param.index_param.indexed_dims) if indexed) assert next(bat_idx_shapes_, None) is None bat_val = draw(hnp.arrays(np.float32, vmap_index_param.bat_slice_shape, elements=_float_elements)) return SetVmapParams(vmap_index_param, bat_ref, bat_val, bat_idxs) Indexer = tuple[int | slice | np.ndarray, ...] def _unpack_idx(idx: Indexer ) -> tuple[Sequence[int | np.ndarray], Sequence[bool]]: indexed_dims = [type(i) != slice for i in idx] non_slice_idx = [i for i, b in zip(idx, indexed_dims) if b] return non_slice_idx, indexed_dims def _pack_idx(non_slice_idx: Sequence[int | np.ndarray], indexed_dims: Sequence[bool]) -> Indexer: idx_ = iter(non_slice_idx) idx = tuple(next(idx_) if b else slice(None) for b in indexed_dims) assert next(idx_, None) is None return idx @jtu.thread_unsafe_test_class(condition=not htu.hypothesis_is_thread_safe()) class StateHypothesisTest(jtu.JaxTestCase): @hp.given(get_vmap_params()) @hp.settings(deadline=None, print_blob=True, max_examples=jtu.NUM_GENERATED_CASES.value, suppress_health_check=[ *hp.settings.default.suppress_health_check, hp.HealthCheck.too_slow, ]) def test_get_vmap(self, get_vmap_param: GetVmapParams): indexed_dims = get_vmap_param.vmap_index_param.index_param.indexed_dims def f(ref, *non_slice_idx): idx = _pack_idx(non_slice_idx, indexed_dims) return [ref_get(ref, idx)] ref_aval = get_vmap_param.vmap_index_param.index_param.ref_aval bat_ref_aval = get_vmap_param.vmap_index_param.bat_ref_aval bat_non_slice_idx_avals = get_vmap_param.vmap_index_param.bat_non_slice_idx_avals ref_bdim = get_vmap_param.vmap_index_param.ref_bdim idx_bdims = get_vmap_param.vmap_index_param.non_slice_idx_bdims out_bdim = get_vmap_param.vmap_index_param.slice_bdim non_slice_idx = get_vmap_param.bat_idxs idx_avals = get_vmap_param.vmap_index_param.index_param.idx_avals ref = get_vmap_param.bat_ref f_batched = jax.vmap(f, in_axes=(ref_bdim, *idx_bdims), out_axes=[out_bdim]) stateful_jaxpr, _ = trace_to_jaxpr(f_batched, bat_ref_aval, *bat_non_slice_idx_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) discharge_of_vmap_ans = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, ref, *non_slice_idx ) # vmap-of-discharge stateful_jaxpr, _ = trace_to_jaxpr(f, ref_aval, *idx_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) f_batched = jax.vmap( partial( core.eval_jaxpr, discharged_jaxpr, discharged_jaxpr.consts ), in_axes=(ref_bdim, *idx_bdims), out_axes=[out_bdim, ref_bdim], ) vmap_of_discharge_ans = f_batched(ref, *non_slice_idx) self.assertAllClose(discharge_of_vmap_ans, vmap_of_discharge_ans, check_dtypes=False) @hp.given(set_vmap_params()) @hp.settings(deadline=None, print_blob=True, max_examples=jtu.NUM_GENERATED_CASES.value, suppress_health_check=[ *hp.settings.default.suppress_health_check, hp.HealthCheck.too_slow, ]) def test_set_vmap(self, set_vmap_param: SetVmapParams): if jtu.test_device_matches(["gpu"]): self.skipTest("Scatter is nondeterministic on GPU") indexed_dims = set_vmap_param.vmap_index_param.index_param.indexed_dims def f(ref, val, *non_slice_idx): idx = _pack_idx(non_slice_idx, indexed_dims) ref_set(ref, idx, val) return [] ref_aval = set_vmap_param.vmap_index_param.index_param.ref_aval bat_ref_aval = set_vmap_param.vmap_index_param.bat_ref_aval bat_non_slice_idx_avals = set_vmap_param.vmap_index_param.bat_non_slice_idx_avals ref_bdim = set_vmap_param.vmap_index_param.ref_bdim idx_bdims = set_vmap_param.vmap_index_param.non_slice_idx_bdims non_slice_idx = set_vmap_param.bat_idxs idx_avals = set_vmap_param.vmap_index_param.index_param.idx_avals ref = set_vmap_param.bat_ref val = set_vmap_param.bat_val bat_val_aval = set_vmap_param.vmap_index_param.bat_slice_aval val_aval = set_vmap_param.vmap_index_param.index_param.slice_aval val_bdim = set_vmap_param.vmap_index_param.slice_bdim f_batched = jax.vmap(f, in_axes=(ref_bdim, val_bdim, *idx_bdims), out_axes=[]) stateful_jaxpr, _ = trace_to_jaxpr( f_batched, bat_ref_aval, bat_val_aval, *bat_non_slice_idx_avals ) discharged_jaxpr = discharge_state(stateful_jaxpr) discharge_of_vmap_ans = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, ref, val, *non_slice_idx, ) # vmap-of-discharge stateful_jaxpr, _ = trace_to_jaxpr(f, ref_aval, val_aval, *idx_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) f_batched = jax.vmap( partial( core.eval_jaxpr, discharged_jaxpr, discharged_jaxpr.consts ), in_axes=(ref_bdim, val_bdim, *idx_bdims), out_axes=[ref_bdim], ) vmap_of_discharge_ans = f_batched(ref, val, *non_slice_idx) self.assertAllClose(discharge_of_vmap_ans, vmap_of_discharge_ans, check_dtypes=False) @hp.given(set_vmap_params()) @hp.settings(deadline=None, print_blob=True, max_examples=jtu.NUM_GENERATED_CASES.value, suppress_health_check=[ *hp.settings.default.suppress_health_check, hp.HealthCheck.too_slow, ]) def test_addupdate_vmap(self, set_vmap_param: SetVmapParams): indexed_dims = set_vmap_param.vmap_index_param.index_param.indexed_dims def f(ref, val, *non_slice_idx): idx = _pack_idx(non_slice_idx, indexed_dims) ref_addupdate(ref, idx, val) return [] ref_aval = set_vmap_param.vmap_index_param.index_param.ref_aval bat_ref_aval = set_vmap_param.vmap_index_param.bat_ref_aval bat_non_slice_idx_avals = set_vmap_param.vmap_index_param.bat_non_slice_idx_avals ref_bdim = set_vmap_param.vmap_index_param.ref_bdim idx_bdims = set_vmap_param.vmap_index_param.non_slice_idx_bdims non_slice_idx = set_vmap_param.bat_idxs idx_avals = set_vmap_param.vmap_index_param.index_param.idx_avals ref = set_vmap_param.bat_ref val = set_vmap_param.bat_val bat_val_aval = set_vmap_param.vmap_index_param.bat_slice_aval val_aval = set_vmap_param.vmap_index_param.index_param.slice_aval val_bdim = set_vmap_param.vmap_index_param.slice_bdim f_batched = jax.vmap(f, in_axes=(ref_bdim, val_bdim, *idx_bdims), out_axes=[]) stateful_jaxpr, _ = trace_to_jaxpr( f_batched, bat_ref_aval, bat_val_aval, *bat_non_slice_idx_avals ) discharged_jaxpr = discharge_state(stateful_jaxpr) discharge_of_vmap_ans = core.eval_jaxpr( discharged_jaxpr, discharged_jaxpr.consts, ref, val, *non_slice_idx, ) # vmap-of-discharge stateful_jaxpr, _ = trace_to_jaxpr(f, ref_aval, val_aval, *idx_avals) discharged_jaxpr = discharge_state(stateful_jaxpr) f_batched = jax.vmap( partial( core.eval_jaxpr, discharged_jaxpr, discharged_jaxpr.consts ), in_axes=(ref_bdim, val_bdim, *idx_bdims), out_axes=[ref_bdim], ) vmap_of_discharge_ans = f_batched(ref, val, *non_slice_idx) self.assertAllClose(discharge_of_vmap_ans, vmap_of_discharge_ans, check_dtypes=False, atol=1e-4, rtol=1e-4) class StateControlFlowTest(jtu.JaxTestCase): def test_cond_discharge(self): def f0(pred, x_ref, y_ref): def true_fun(): x_ref[...] = 1. def false_fun(): y_ref[...] = 2. lax.cond(pred, true_fun, false_fun) return x_ref[...], y_ref[...] ref = lambda x: AbstractRef(core.typeof(x)) f_jaxpr = jax.make_jaxpr(f0)(False, ref(3.), ref(4.)) jaxpr = discharge_state(f_jaxpr, should_discharge=[False, False, True]) # Effects on y_ref were discharged away but not the effects on x_ref _, x, y = f_jaxpr.invars self.assertEqual(f_jaxpr.effects, {ReadEffect(x), WriteEffect(x), ReadEffect(y), WriteEffect(y)}) self.assertEqual(jaxpr.effects, {ReadEffect(jaxpr.invars[1]), WriteEffect(jaxpr.invars[1])}) # x_ref arg is still a reference but y_ref is discharged self.assertNotIsInstance(jaxpr.invars[2].aval, AbstractRef) self.assertIsInstance(jaxpr.invars[1].aval, AbstractRef) # x_ref value is returned as part of the discharged refs set. self.assertLen(f_jaxpr.out_avals, 2) self.assertLen(jaxpr.outvars, 3) def test_cond_with_ref_reuse(self): def f(pred): def body(x_ref): def true_fun(): x_ref[()] = 1. def false_fun(): x_ref[()] = 2. lax.cond(pred, true_fun, false_fun) return run_state(body)(0.) jaxpr = jax.make_jaxpr(f)(True).jaxpr self.assertEmpty(jaxpr.effects) out_true = jax.jit(f)(True) expected_true = 1. self.assertAllClose(out_true, expected_true) out_false = jax.jit(f)(False) expected_false = 2. self.assertAllClose(out_false, expected_false) def test_cond_readonly_refs(self): def f(pred): def body(refs): x_ref, y_ref, z_ref = refs def true_fun(): y_ref[()] = x_ref[()] def false_fun(): y_ref[()] = x_ref[()] + z_ref[()] lax.cond(pred, true_fun, false_fun) return run_state(body)((1., 0., 2.)) jaxpr = jax.make_jaxpr(f)(True).jaxpr [run_state_eqn] = jaxpr.eqns *_, cond_eqn = discharge_state(run_state_eqn.params["jaxpr"]).eqns self.assertIs(cond_eqn.primitive, lax.cond_p) self.assertLen(cond_eqn.invars, 4) # pred + 3x ref values self.assertLen(cond_eqn.outvars, 1) # only the updated ref value self.assertAllClose(jax.jit(f)(True), (1., 1., 2.)) self.assertAllClose(jax.jit(f)(False), (1., 3., 2.)) def test_simple_cond_using_multiple_refs_with_interleaved_consts(self): def f(pred): def body(refs): x_ref, y_ref, z_ref = refs a = jnp.array(2.) b = jnp.array(2.) def true_fun(): x_ref[()] = 1. * a y_ref[()] = 1. def false_fun(): z_ref[()] = 1. x_ref[()] = 2. x_ref[()] = 2. * b lax.cond(pred, true_fun, false_fun) return run_state(body)((0., 0., 0.)) jaxpr = jax.make_jaxpr(f)(True).jaxpr self.assertEmpty(jaxpr.effects) out_true = jax.jit(f)(True) expected_true = (2., 1., 0.) self.assertAllClose(out_true, expected_true) out_false = jax.jit(f)(False) expected_false = (4., 0., 1.) self.assertAllClose(out_false, expected_false) def test_cond_nested_run_state_using_multiple_refs_with_interleaved_consts( self): def f(pred): @run_state def body(refs): x_ref, o_ref1, o_ref2 = refs @run_state def body2(refs): y_ref, o_ref3 = refs @run_state def body3(z_ref): a = jnp.array(2.) b = jnp.array(2.) def true_fun(): x_ref[()] = 1. * a y_ref[()] = 1. def false_fun(): z_ref[()] = 1. x_ref[()] = 2. x_ref[()] = 2. * b lax.cond(pred, true_fun, false_fun) o_ref3[...] = body3(0.) o_ref1[...], o_ref2[...] = body2((0., 0.)) return body((0., 0., 0.)) jaxpr = jax.make_jaxpr(f)(True).jaxpr self.assertEmpty(jaxpr.effects) out_true = jax.jit(f)(True) expected_true = (2., 1., 0.) self.assertAllClose(out_true, expected_true) out_false = jax.jit(f)(False) expected_false = (4., 0., 1.) self.assertAllClose(out_false, expected_false) def test_nested_cond_using_multiple_refs_with_interleaved_consts(self): def f(pred1, pred2): def body(refs): x_ref, y_ref, z_ref = refs a = jnp.array(2.) def true_fun(): b = jnp.array(2.) def inner_true_fun(): x_ref[()] = 1. * a y_ref[()] = 1. def inner_false_fun(): z_ref[()] = 1. x_ref[()] = 2. x_ref[()] = 2. * b lax.cond(pred2, inner_true_fun, inner_false_fun) def false_fun(): a = jnp.array(2.) b = jnp.array(2.) def inner_true_fun(): x_ref[()] = 1. * a z_ref[()] = 4. def inner_false_fun(): x_ref[()] = 2. * b lax.cond(pred2, inner_true_fun, inner_false_fun) lax.cond(pred1, true_fun, false_fun) return run_state(body)((0., 0., 0.)) jaxpr = jax.make_jaxpr(f)(True, False).jaxpr self.assertEmpty(jaxpr.effects) out = jax.jit(f)(True, True) expected = (2., 1., 0.) self.assertTupleEqual(out, expected) out = jax.jit(f)(True, False) expected = (4., 0., 1.) self.assertTupleEqual(out, expected) out = jax.jit(f)(False, True) expected = (2., 0., 4.) self.assertTupleEqual(out, expected) out = jax.jit(f)(False, False) expected = (4., 0., 0.) self.assertTupleEqual(out, expected) def test_while_with_state_in_body(self): def f(x, y, z): @run_state def body(x_ref): def cond(i): return i < y def body(i): x_ref[...] += z return i + 1 lax.while_loop(cond, body, 0) return body(x) jaxpr = jax.make_jaxpr(f)(0, 5, 2).jaxpr self.assertEmpty(jaxpr.effects) self.assertAllClose(jax.jit(f)(0, 5, 2), 10) self.assertAllClose(jax.jit(f)(1, 2, 3), 7) def test_while_with_state_in_cond(self): def f(x, y, z): @run_state def body(x_ref): def cond(i): x_ref[...] += z return i < y def body(i): return i + 1 lax.while_loop(cond, body, 0) return body(x) jaxpr = jax.make_jaxpr(f)(0, 5, 2).jaxpr self.assertEmpty(jaxpr.effects) self.assertAllClose(jax.jit(f)(0, 5, 2), 10) self.assertAllClose(jax.jit(f)(1, 2, 3), 7) def test_while_errors_if_same_ref_in_body_and_cond(self): def f(x, y, z): @run_state def body(x_ref): def cond(i): x_ref[...] += z return i < y def body(i): x_ref[...] += z return i + 1 lax.while_loop(cond, body, 0) return body(x) jaxpr = jax.make_jaxpr(f)(0, 5, 2).jaxpr self.assertEmpty(jaxpr.effects) with self.assertRaisesRegex(NotImplementedError, "Cannot write to the same ref in both cond and body."): jax.jit(f)(0, 5, 2) def test_scan_with_state_in_body(self): def f(x, w, y, zs): @run_state def body(refs): x_ref, w_ref = refs def body(y, z): x_ref[...] += y w_ref[...] += z return y + 1, () lax.scan(body, y, zs) return body((x, w)) zs = jnp.arange(5) jaxpr = jax.make_jaxpr(f)(0, 1, 5, zs).jaxpr self.assertEmpty(jaxpr.effects) self.assertAllClose(jax.jit(f)(0, 1, 5, zs), (35, 11)) self.assertAllClose(jax.jit(f)(1, 1, 2, zs), (21, 11)) def test_scan_with_state_in_body_nested(self): @run_state def g(refs): a_ref, x_ref, w_ref, y_ref, zs_ref = refs def f(x, w, y, zs): @run_state def loop(refs): x_ref, w_ref = refs def body(y, z): x_ref[...] += y w_ref[...] += z a_ref[...] += z return y + 1, () lax.scan(body, y, zs) a_ref[...] += 1 out = loop((x, w)) a_ref[...] += 1 return out x, w = f(x_ref[...], w_ref[...], y_ref[...], zs_ref[...]) x_ref[...] = x w_ref[...] = w zs = jnp.arange(5) jaxpr = jax.make_jaxpr(g)((1, 0, 1, 5, zs)).jaxpr self.assertEmpty(jaxpr.effects) self.assertAllClose(jax.jit(g)((1, 0, 1, 5, zs))[:3], (13, 35, 11)) self.assertAllClose(jax.jit(g)((1, 1, 1, 2, zs))[:3], (13, 21, 11)) def test_scan_discharges_into_carry(self): # we want to discharge scanned-over refs into the carry for aliasing def body(_, x_ref): x_ref[...] += 1 return (), () x_ref = jax.new_ref(jnp.arange(3.)) jaxpr = jax.make_jaxpr(lambda x_ref: jax.lax.scan(body, (), x_ref))(x_ref) jaxpr = discharge_state(jaxpr) scan_eqn = jaxpr.eqns[0] consts_g, carry_g, _ = scan_eqn.params['ft_in'].unpack() self.assertEqual([len(consts_g), len(carry_g)], [0, 2]) a, b = scan_eqn.params['jaxpr'].in_avals self.assertEqual(a.shape, ()) self.assertEqual(b.shape, (3,)) @parameterized.named_parameters( ("call_primitive", core.call_p), ("closed_call_primitive", core.closed_call_p), ) def test_call_primitive_discharges(self, prim): def g(y_ref, x): x_ref = jax.new_ref(x) y_ref[...] = jnp.exp(x_ref[...]) return [jax.freeze(y_ref)] def f(x): y_ref = jax.new_ref(jnp.zeros_like(x)) g_ = partial(g, y_ref) sub = lu.wrap_init(g_, debug_info=api_util.debug_info("f", g, (x,), {})) jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(sub, [core.typeof(x)]) return prim.bind(*consts, x, call_jaxpr=jaxpr)[0] out = f(4.) np.testing.assert_array_equal(out, jnp.exp(4.)) def test_checkpoint_internal_ref_mutation_grad(self): x = jnp.ones((4, 4)) def f(x): ref = jax.new_ref(jnp.zeros_like(x)) ref[...] = x + 1.0 return jax.ref.freeze(ref) with config.remat3(False), self.assertRaisesRegex( NotImplementedError, "Mutating a reference with unknown/differentiated values", ): jax.grad(lambda x: jnp.sum(jax.checkpoint(f)(x)))(x) with config.remat3(True): grad = jax.grad(lambda x: jnp.sum(jax.checkpoint(f)(x)))(x) np.testing.assert_allclose(grad, jnp.ones((4, 4))) def test_transformed_ref_is_a_pytree_node(self): @jax.jit def fn(): x_ref = jax.new_ref(jnp.zeros(7)) x_ref_view = x_ref.at[1:6] def assert_fn(x): # ensure TransformedRef is still a node self.assertNotIsInstance(x, jax.Array) self.assertIsInstance(x, state_types.TransformedRef) jax.tree.map(assert_fn, x_ref_view) fn() def test_transformed_ref_memory_space(self): @jax.jit def fn(i): x_ref = jax.new_ref(jnp.zeros((7, 8)), memory_space=core.MemorySpace.Device) self.assertEqual(x_ref.at[1:6].memory_space, core.MemorySpace.Device) self.assertEqual(x_ref.at[i].memory_space, core.MemorySpace.Device) self.assertEqual(x_ref.at[i, 0].memory_space, core.MemorySpace.Device) fn(0) def test_tracer_memory_space_raises_value_error(self): @jax.jit def f(x): return x.memory_space with self.assertRaisesRegex( ValueError, "The 'memory_space' attribute is not available on" ): f(jnp.zeros((2, 3))) def test_transformed_multiref_memory_space(self): @jax.jit def fn(idx): x_ref = jax.new_ref(jnp.zeros((7, 8)), memory_space=core.MemorySpace.Device) y_ref = jax.new_ref(jnp.zeros((7, 8)), memory_space=core.MemorySpace.Device) multi_ref = state_types.TransformedRef( ref=(x_ref, y_ref), transforms=(state_types.SelectTransform(idx),), ) self.assertEqual(multi_ref.memory_space, core.MemorySpace.Device) fn(0) def test_transformed_multiref_inconsistent_memory_space_error(self): @jax.jit def fn(idx): x_ref = jax.new_ref(jnp.zeros((7, 8)), memory_space=core.MemorySpace.Device) y_ref = jax.new_ref(jnp.zeros((7, 8)), memory_space=None) multi_ref = state_types.TransformedRef( ref=(x_ref, y_ref), transforms=(state_types.SelectTransform(idx),), ) with self.assertRaisesRegex( TypeError, "Cannot select from Refs of different types" ): _ = multi_ref.memory_space fn(0) @parameterized.named_parameters( ("scan", "scan"), ("while_loop", "while_loop")) def test_transformed_ref_in(self, transform): def scan_by_closing_over_fn(fn): return lambda *args: jax.lax.scan(lambda _, __: (fn(*args), ()), None, (), length=1)[0] def while_loop_by_closing_over(fn): return lambda *args: jax.lax.while_loop( lambda i: i < 1, lambda i: (fn(*args), i + 1)[1], 0) transform = {"scan": scan_by_closing_over_fn, "while_loop": while_loop_by_closing_over}[transform] @jax.jit def f(): x_ref = jax.new_ref(jnp.zeros(7)) x_ref_view = x_ref.at[1:6] @transform def inner(x_ref_view): x_ref_view[...] += 17 * jnp.ones((5,)) inner(x_ref_view) return x_ref[...] expected = jnp.zeros(7).at[1:6].add(17.0) self.assertAllClose(f(), expected) @parameterized.named_parameters(("vmap", "vmap")) def test_no_transformed_ref_in(self, transform): def scan_fn(fn): return lambda *args: jax.lax.scan( lambda _, ref_slices_maybe: (fn(*ref_slices_maybe), ()), None, args, length=1)[0] def while_loop_fn(fn): return lambda *args: jax.lax.while_loop( lambda _: False, lambda ref_slices_maybe: fn(*ref_slices_maybe), args) def cond_fn(fn): return lambda *args: jax.lax.cond(True, fn, fn, *args) def shard_map(fn): mesh = jax.make_mesh((1,), ("x",), axis_types=jax.sharding.AxisType.Explicit) def fn_(*args): return jax.shard_map( fn, mesh=mesh, in_specs=jax.P(), out_specs=jax.P())(*args) return fn_ transform = {"jit": jax.jit, "vmap": jax.vmap, "remat": jax.remat, "scan": scan_fn, "while_loop": while_loop_fn, "cond": cond_fn, "shard_map": shard_map}[transform] @jax.jit def f(): x_ref = jax.new_ref(jnp.zeros(7)) x_ref_view = x_ref.at[1:6] @transform def inner(x_ref_view): x_ref_view[...] += jnp.ones_like(x_ref_view) inner(x_ref_view) return x_ref[...] with self.assertRaisesRegex(TypeError, "TransformedRefs are not allowed"): f() @parameterized.named_parameters(("jit", "jit"), ("remat", "remat"), ("shard_map", "shard_map")) def test_transformed_ref_in_jit_remat_shard_map(self, transform): transform = {"jit": jax.jit, "remat": jax.remat, "shard_map": partial(_shard_map, out_specs=None)}[transform] @jax.jit def f(): x_ref = jax.new_ref(jnp.zeros(7)) x_ref_view = x_ref.at[1:6] @transform def inner(x_ref_view): x_ref_view[...] += 17 * jnp.ones((5,)) inner(x_ref_view) return x_ref[...] expected = jnp.zeros(7).at[1:6].add(17.0) self.assertAllClose(f(), expected) def test_transformed_ref_in_cond(self): @jax.jit def f(pred): x_ref = jax.new_ref(jnp.zeros(7)) x_ref_view = x_ref.at[1:6] def true_fn(x_ref_view): x_ref_view[...] += 17 * jnp.ones((5,)) def false_fn(x_ref_view): x_ref_view[...] += 5 * jnp.ones((5,)) jax.lax.cond(pred, true_fn, false_fn, x_ref_view) return x_ref[...] expected_true = jnp.zeros(7).at[1:6].add(17.0) expected_false = jnp.zeros(7).at[1:6].add(5.0) self.assertAllClose(f(True), expected_true) self.assertAllClose(f(False), expected_false) @parameterized.named_parameters( ("remat", jax.remat), ("cond", _cond), ("jit", jax.jit), ("shard_map", partial(_shard_map, out_specs=P()))) def test_read_transformed_ref(self, transform): @jax.jit def f(): x_ref = jax.new_ref(jnp.arange(10, dtype=jnp.float32)) x_view = x_ref.at[2:7] result = transform(lambda v: v[...])(x_view) return result expected = jnp.arange(10, dtype=jnp.float32)[2:7] self.assertAllClose(f(), expected) @parameterized.named_parameters(("false", False), ("true", True)) def test_cond_write_transformed_or_pure_ref(self, pure_ref): x_ref = jax.new_ref(jnp.zeros(10, dtype=jnp.float32)) @jax.jit def f(x_ref): x_view = x_ref.at[2:7] def body(v): if pure_ref: v[2:7] = jnp.ones_like(v[2:7]) else: v[...] = jnp.ones_like(v) if pure_ref: _cond(body)(x_ref) else: _cond(body)(x_view) return x_ref[...] expected = jnp.zeros(10, dtype=x_ref.dtype).at[2:7].set(1.0) self.assertAllClose(f(x_ref), expected) def test_transformed_ref_in_jit_simple(self): @jax.jit def f(): x_ref = jax.new_ref(jnp.zeros(3)) # x_ref: Ref{f32[3]} x_ref_view = x_ref.at[1] inner({"x_ref_view": x_ref_view, "dummy": 1}) @jax.jit def inner(x_ref_view_dict): x_ref_view = x_ref_view_dict["x_ref_view"] x_ref_view[...] = jnp.ones_like(x_ref_view) print(f.trace().jaxpr) @parameterized.named_parameters( ("remat", jax.remat), ("cond", _cond), ("jit", jax.jit), ("shard_map", partial(_shard_map, out_specs=None))) def test_write_transformed_ref(self, transform): @jax.jit def f(): x_ref = jax.new_ref(jnp.zeros(10, dtype=jnp.float32)) x_view = x_ref.at[2:7] def body(v): v[...] = jnp.ones_like(v) transform(body)(x_view) return x_ref[...] expected = jnp.zeros(10, dtype=jnp.float32).at[2:7].set(1.0) self.assertAllClose(f(), expected) @parameterized.named_parameters( ("remat", jax.remat), ("cond", _cond), ("jit", jax.jit), ("shard_map", partial(_shard_map, out_specs=None))) def test_addupdate_transformed_ref(self, transform): @jax.jit def f(): x_ref = jax.new_ref(jnp.ones(10, dtype=jnp.float32)) x_view = x_ref.at[2:7] def body(v): v[...] += 1.0 transform(body)(x_view) return x_ref[...] expected = jnp.ones(10, dtype=jnp.float32).at[2:7].add(1.0) self.assertAllClose(f(), expected) class GeneralRefTest(jtu.JaxTestCase): def test_token(self): def f(x_ref): x = x_ref[...] x_ref[...] = x ref_addupdate(x_ref, (), x) return [x] jaxpr, _ = trace_to_jaxpr(f, AbstractRef(core.AbstractToken())) self.assertIs(type(jaxpr.outvars[0].aval), core.AbstractToken) def test_reshape(self): def f(x_ref): x_ref = x_ref.reshape(4, -1) x_ref.reshape(-1)[...] = jnp.arange(36) return [x_ref[...]] jaxpr, _ = trace_to_jaxpr( f, AbstractRef(core.ShapedArray((12, 3), jnp.int32)) ) self.assertEqual(jaxpr.outvars[0].aval.shape, (4, 9)) # NOTE(mattjj): disabled because it's extremely illegal # def test_ref_of_ref(self): # def f(x_ref_ref): # x_ref = x_ref_ref[...] # return [x_ref] # # Not sure why you'd ever want to do this, but it works! # jaxpr, _, _ = pe.trace_to_jaxpr_dynamic( # wrap_init(f, 1), # [AbstractRef(AbstractRef(core.ShapedArray((), jnp.int32)))]) # self.assertIs(type(jaxpr.outvars[0].aval), AbstractRef) # self.assertIs(type(jaxpr.outvars[0].aval.inner_aval), core.ShapedArray) class RunStateTest(jtu.JaxTestCase): def test_simple_run_state(self): out = run_state(lambda _: None)(1) self.assertEqual(out, 1) def test_nontrivial_run_state(self): def f(refs): x_ref, y_ref = refs x = x_ref[...] * y_ref[...] y_ref[...] = x * 2 x_ref[...] = y_ref[...] + x_ref[...] # x + x * y * 2, x * y * 2 x, y = run_state(f)((2, 3)) self.assertEqual(x, 2 + 2 * 3 * 2) self.assertEqual(y, 2 * 3 * 2) def test_run_state_with_uninitialized_input(self): def f(refs): x_ref, y_ref = refs # y_ref is uninitialized so we shouldn't read from it until we write into # it. x = x_ref[...] y_ref[...] = x * 2 x_ref[...] = y_ref[...] + x_ref[...] # x + x * 2, x * 2 # jax.ShapeDtypeStruct is weirdly special to JAX, so we make our own class. class MyArrayType: pass state_types._ref_type_aval_mappings[MyArrayType] = lambda _: ( AbstractRef(core.ShapedArray((), jnp.int32)), state_types.uninitialized, ) x, y = run_state(f)((jnp.int32(2), MyArrayType())) self.assertEqual(x, 2 + 2 * 2) self.assertEqual(y, 2 * 2) def test_nontrivial_run_state_jit(self): def f(refs): x_ref, y_ref = refs @jax.jit def g(): x = x_ref[...] * y_ref[...] y_ref[...] = x * 2 x_ref[...] = y_ref[...] + x_ref[...] # x + x * y * 2, x * y * 2 g() x, y = run_state(f)((2, 3)) self.assertEqual(x, 2 + 2 * 3 * 2) self.assertEqual(y, 2 * 3 * 2) def test_simple_run_state_with_multiple_refs(self): out1, out2 = run_state(lambda _: None)((1, 2)) self.assertEqual(out1, 1) self.assertEqual(out2, 2) def test_simple_run_state_with_tuple(self): out1, out2 = run_state(lambda _: None)((1, 2)) self.assertEqual(out1, 1) self.assertEqual(out2, 2) def test_can_stage_run_state(self): def f(x): return run_state(lambda _: None)(x) jaxpr = jax.make_jaxpr(f)(2) self.assertIsNotNone(jaxpr.debug_info) self.assertIsNotNone(jaxpr.debug_info.func_src_info) # NOTE(mattjj): disabled because the error message changed for the better # def test_can_stage_run_state_leaked_tracer_error(self): # leaks = [] # def f(x): # def my_fun(x): # leaks.append(x) # return None # return run_state(my_fun)(x) # _ = jax.make_jaxpr(f)(2) # with self.assertRaisesRegex(jax.errors.UnexpectedTracerError, # "The function being traced when the value leaked was .*my_fun"): # jax.jit(lambda _: leaks[0])(1) def test_nested_run_state_captures_effects(self): def f(x): def body(x_ref): def inner(y_ref): y_ref[...] x_ref[...] run_state(inner)(1) return run_state(body)(x) jaxpr = jax.make_jaxpr(f)(2) self.assertEmpty(jaxpr.effects) self.assertEmpty(jaxpr.eqns[0].effects) body_jaxpr = jaxpr.eqns[0].params["jaxpr"] self.assertSetEqual(body_jaxpr.effects, {ReadEffect(body_jaxpr.invars[0])}) inner_jaxpr = body_jaxpr.eqns[0].params["jaxpr"] self.assertSetEqual(inner_jaxpr.effects, {ReadEffect(inner_jaxpr.invars[0]), ReadEffect(inner_jaxpr.invars[1])}) def test_jvp_of_run_state(self): @run_state def f(refs): x_ref, y_ref = refs y_ref[...] = jnp.sin(x_ref[...]) xy, xy_t = jax.jvp(f, ((2., 1.),), ((3., 1.),)) # x, sin(x) self.assertAllClose(xy, (2., np.sin(2.))) # t, cos(x) * t self.assertAllClose(xy_t, (3., 3 * np.cos(2.))) x, x_t = jax.jvp(lambda x: f((x, 0.))[1], (2.,), (3.,)) self.assertAllClose(x, np.sin(2.)) self.assertAllClose(x_t, 3 * np.cos(2.)) def test_jvp_of_run_state_with_zero_tangent(self): @run_state def f(refs): x_ref, z_ref, y_ref = refs del z_ref y_ref[...] = jnp.sin(x_ref[...]) x, x_t = jax.jvp(lambda x: f((x, 0., 0.,))[2], (2.,), (3.,)) self.assertAllClose(x, np.sin(2.)) self.assertAllClose(x_t, 3 * np.cos(2.)) class FuncSpec(NamedTuple): fun: Callable[..., Any] name: str min_rank: int = 0 max_rank: int = 4 min_dim: int = 0 max_dim: int = 4 def call(self, *args): return run_state(self.fun)(*args) def ref(self, *args): return run_state_reference(self.fun)(*args) def sin_stateful(refs): x_ref, y_ref = refs y_ref[...] = jnp.sin(x_ref[...]) sin_spec = FuncSpec(sin_stateful, "sin") def cos_stateful(refs): x_ref, y_ref = refs y_ref[...] = jnp.cos(x_ref[...]) cos_spec = FuncSpec(cos_stateful, "cos") def mul2_stateful(refs): x_ref, y_ref = refs y_ref[...] = x_ref[...] y_ref[...] = y_ref[...] + x_ref[...] mul2_spec = FuncSpec(mul2_stateful, "mul2") def mul2_stateful_with_constant(refs): x_ref, y_ref = refs y_ref[...] = (2. * np.ones(x_ref.shape, x_ref.dtype)) * x_ref[...] mul2_constant_spec = FuncSpec(mul2_stateful_with_constant, "mul2_c") def crazy_identity_stateful(refs): x_ref, y_ref = refs x = x_ref[...] x_ref[...] = (x + x) / 2 y_ref[...] = x_ref[...] y = y_ref[...] y_ref[...] = (y + y) / 2 crazy_identity_spec = FuncSpec(crazy_identity_stateful, "id") def func_spec(depth: int = 4): raw_specs = hps.sampled_from([sin_spec, cos_spec, mul2_spec, mul2_constant_spec, crazy_identity_spec]) if depth > 0: return hps.one_of([raw_specs, nest_spec(depth - 1), add_spec(depth - 1), compose_spec(depth - 1)]) return raw_specs @hps.composite def compose_spec(draw, depth): f1 = draw(func_spec(depth)) f2 = draw(func_spec(depth)) def wrapped_impl(*args): f1.fun(*args) f2.fun(*args) return FuncSpec(wrapped_impl, f"({f2.name} . {f1.name})", min_rank=max(f1.min_rank, f2.min_rank), max_rank=min(f1.max_rank, f2.max_rank), min_dim=max(f1.min_dim, f2.min_dim), max_dim=min(f1.max_dim, f2.max_dim)) @hps.composite def nest_spec(draw, depth): f = draw(func_spec(depth)) def wrapped_impl(refs): x_ref, y_ref = refs x, y = x_ref[...], y_ref[...] x, y = run_state(f.fun)((x, y)) x_ref[...], y_ref[...] = x, y return FuncSpec(wrapped_impl, f"nest({f.name})", min_rank=f.min_rank, max_rank=f.max_rank, min_dim=f.min_dim, max_dim=f.max_dim) @hps.composite def add_spec(draw, depth): f1 = draw(func_spec(depth)) f2 = draw(func_spec(depth)) def wrapped_impl(refs): x_ref, y_ref = refs x, y = x_ref[...], y_ref[...] x1, y1 = run_state(f1.fun)((x, y)) x2, y2 = run_state(f2.fun)((x, y)) x_ref[...], y_ref[...] = x1 + x2, y1 + y2 return FuncSpec(wrapped_impl, f"({f2.name} + {f1.name})", min_rank=max(f1.min_rank, f2.min_rank), max_rank=min(f1.max_rank, f2.max_rank), min_dim=max(f1.min_dim, f2.min_dim), max_dim=min(f1.max_dim, f2.max_dim)) @jtu.thread_unsafe_test_class(condition=not htu.hypothesis_is_thread_safe()) class RunStateHypothesisTest(jtu.JaxTestCase): @jax.legacy_prng_key('allow') @hp.given(hps.data()) @hp.settings(deadline=None, print_blob=True, max_examples=jtu.NUM_GENERATED_CASES.value) def test_jvp(self, data): spec = data.draw(func_spec()) def impl(x): return spec.call((x, jnp.zeros_like(x)))[1] def ref(x): return spec.ref((x, jnp.zeros_like(x)))[1] k1, k2 = random.split(random.PRNGKey(0)) shape = data.draw(hnp.array_shapes(min_dims=spec.min_rank, max_dims=spec.max_rank, min_side=spec.min_dim, max_side=spec.max_dim)) x = random.normal(k1, shape) t = random.normal(k2, x.shape) y, y_t = jax.jvp(impl, (x,), (t,)) y_ref, y_ref_t = jax.jvp(ref, (x,), (t,)) self.assertAllClose(y, y_ref) self.assertAllClose(y_t, y_ref_t) class PinnedBuffersTest(jtu.JaxTestCase): def test_pin_unpin_basic(self): @jax.jit def f(x): return unpin(pin(x)) x = jnp.arange(3.) txt = f.lower(x).as_text('hlo') self.assertIn("Pin", txt) if jtu.test_device_matches(['gpu', 'tpu']): y = f(x) self.assertAllClose(y, x) def test_new_ref_pinned(self): @jax.jit def f(x): ref = jax.new_ref(x, pin=True) return jax.ref.freeze(ref) x = jnp.arange(3.) txt = f.lower(x).as_text('hlo') self.assertIn("Pin", txt) self.assertIn("Unpin", txt) if jtu.test_device_matches(['gpu', 'tpu']): y = f(x) self.assertAllClose(y, x) def test_new_ref_unpinned_by_default(self): @jax.jit def f(x): ref = jax.new_ref(x) ref[...] = ref[...] + 1. return jax.ref.freeze(ref) x = jnp.arange(3.) txt = f.lower(x).as_text('hlo') self.assertNotIn("Pin", txt) y = f(x) self.assertAllClose(y, x + 1.) def test_empty_ref_pinned(self): @jax.jit def f(): ref = jax.empty_ref(jax.ShapeDtypeStruct((3,), jnp.float32), pin=True) return jax.ref.freeze(ref) txt = f.lower().as_text('hlo') self.assertIn("Pin", txt) self.assertIn("Unpin", txt) if jtu.test_device_matches(['gpu', 'tpu']): y = f() # contents are uninitialized, only check the type self.assertEqual(y.shape, (3,)) self.assertEqual(y.dtype, jnp.float32) def test_free_ref_pinned(self): def f(x): ref = jax.new_ref(x, pin=True) jax.ref.free_ref(ref) return x + 1. x = jnp.arange(3.) jaxpr = jax.make_jaxpr(f)(x) discharged = discharge_state(jaxpr) prims = [e.primitive.name for e in discharged.eqns] self.assertIn('pin', prims) self.assertIn('unpin', prims) if jtu.test_device_matches(['gpu', 'tpu']): y = jax.jit(f)(x) self.assertAllClose(y, x + 1.) def test_new_ref_pinned_eager_error(self): with self.assertRaises(NotImplementedError): jax.new_ref(jnp.arange(3.), pin=True) def test_discharged_aval_non_ref_preserves_memory_space(self): aval = core.ShapedArray((3,), jnp.float32, memory_space=core.MemorySpace.Host) discharged = discharged_aval(aval, discharge=True, strip_memory_space=True) self.assertEqual(discharged.memory_space, core.MemorySpace.Host) self.assertEqual(discharged, aval) discharged_no_strip = discharged_aval( aval, discharge=False, strip_memory_space=True ) self.assertEqual(discharged_no_strip, aval) def test_discharge_state_preserves_pure_input_memory_space(self): def f(pure_x, ref_init): ref = jax.new_ref(ref_init) ref[...] += 1.0 return pure_x, ref[...] x_aval = core.ShapedArray( (4,), jnp.float32, memory_space=core.MemorySpace.Host ) ref_init_aval = core.ShapedArray((4,), jnp.float32) closed_jaxpr = jax.make_jaxpr(f)(x_aval, ref_init_aval) discharged_jaxpr = discharge_state(closed_jaxpr, strip_memory_space=True) self.assertEqual( discharged_jaxpr.jaxpr.invars[0].aval.memory_space, core.MemorySpace.Host, ) def test_shard_map_ref_preserves_memory_space_and_kind(self): mesh = jtu.create_mesh((1,), ("x",)) ref_aval = AbstractRef( core.ShapedArray((4,), jnp.int32), memory_space=core.MemorySpace.Host, kind="custom_kind", ) sharded = core.shard_aval( mesh, frozenset({"x"}), False, jax.sharding.PartitionSpec(), ref_aval ) self.assertIsInstance(sharded, AbstractRef) self.assertEqual(sharded.memory_space, core.MemorySpace.Host) self.assertEqual(sharded.kind, "custom_kind") with jax.set_mesh(mesh): def f(x): host_ref = jax.new_ref(x, memory_space=core.MemorySpace.Host) @jax.shard_map( in_specs=( jax.sharding.PartitionSpec(), jax.sharding.PartitionSpec(), ), out_specs=(), check_vma=False, ) def _write(v, r): self.assertEqual(jax.typeof(r).memory_space, core.MemorySpace.Host) r[...] = v return () _write(jnp.ones((4,), dtype=jnp.int32), host_ref) return host_ref[...] jax.make_jaxpr(f)(jnp.zeros((4,), dtype=jnp.int32)) class WithMemorySpaceConstraintTest(jtu.JaxTestCase): def test_basic_and_jaxpr(self): x = jnp.ones((2, 3), jnp.float32) def f(x): return with_memory_space_constraint(x, core.MemorySpace.Device) jaxpr = jax.make_jaxpr(f)(x) self.assertEqual(jaxpr.out_avals[0].memory_space, core.MemorySpace.Device) np.testing.assert_array_equal(jax.jit(f)(x), x) def test_vmap(self): x = jnp.ones((4, 2, 3), jnp.float32) def f(x): return with_memory_space_constraint(x, core.MemorySpace.Device) out = jax.vmap(f)(x) np.testing.assert_array_equal(out, x) def test_grad(self): def f(x): y = with_memory_space_constraint(x, core.MemorySpace.Device) return jnp.sum(y ** 2) x = jnp.array([2.0, 3.0], jnp.float32) grad = jax.grad(f)(x) np.testing.assert_array_equal(grad, 2.0 * x) class RefTransformTest(jtu.JaxTestCase): def test_swap_reshape(self): @jax.jit def f(): x = jnp.arange(12, dtype=jnp.float32).reshape((4, 3)) ref = jax.new_ref(x) old_val = ref.reshape(6, 2).swap(jnp.zeros((6, 2), dtype=jnp.float32)) return old_val, ref[...] old_val, new_val = f() self.assertEqual(old_val.shape, (6, 2)) self.assertAllClose(old_val, jnp.arange(12, dtype=jnp.float32).reshape((6, 2))) self.assertEqual(new_val.shape, (4, 3)) self.assertAllClose(new_val, jnp.zeros((4, 3), dtype=jnp.float32)) def test_set_reshape(self): @jax.jit def f(): x = jnp.zeros((4, 3)) ref = jax.new_ref(x) ref.reshape(2, 6)[...] = jnp.ones((2, 6)) return ref[...] self.assertAllClose(f(), jnp.ones((4, 3))) def test_bitcast_swap_and_set(self): @jax.jit def f(): ref = jax.new_ref(jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=jnp.float32)) old_val = ref.bitcast(jnp.int32).swap(jnp.zeros((2, 2), dtype=jnp.int32)) ref.bitcast(jnp.int32)[...] = jnp.ones((2, 2), dtype=jnp.int32) return old_val, ref[...] old_val, new_val = f() self.assertEqual(old_val.dtype, jnp.int32) self.assertArraysEqual( old_val, lax.bitcast_convert_type( jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=jnp.float32), jnp.int32 ), ) self.assertEqual(new_val.dtype, jnp.float32) self.assertArraysEqual( new_val, lax.bitcast_convert_type(jnp.ones((2, 2), dtype=jnp.int32), jnp.float32), ) def test_set_reshaped_slice(self): @jax.jit def f(): x = jnp.zeros((4, 3), dtype=jnp.float32) ref = jax.new_ref(x) ref.reshape(2, 6)[0, :3] = jnp.ones(3, dtype=jnp.float32) return ref[...] expected = jnp.array([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=jnp.float32) self.assertAllClose(f(), expected) def test_swap_reshaped_slice(self): @jax.jit def f(): x = jnp.arange(12, dtype=jnp.float32).reshape((4, 3)) ref = jax.new_ref(x) old_val = ref.reshape(2, 6).at[0, :3].swap(jnp.zeros(3, dtype=jnp.float32)) return old_val, ref[...] old_val, new_val = f() self.assertEqual(old_val.shape, (3,)) self.assertAllClose(old_val, jnp.arange(3, dtype=jnp.float32)) expected = jnp.arange(12, dtype=jnp.float32).reshape((4, 3)) expected = expected.reshape(2, 6).at[0, :3].set(0.0).reshape(4, 3) self.assertAllClose(new_val, expected) def test_bitcast_slice_set(self): @jax.jit def f(): ref = jax.new_ref(jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=jnp.float32)) ref.bitcast(jnp.int32)[0, :] = jnp.zeros(2, dtype=jnp.int32) return ref[...] expected = jnp.array([[0.0, 0.0], [3.0, 4.0]], dtype=jnp.float32) self.assertArraysEqual(f(), expected) def test_addupdate_on_reshaped_view(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((4, 3))) ref_addupdate(ref.reshape(2, 6), (), jnp.ones((2, 6))) return ref[...] expected = jnp.ones((4, 3)) self.assertAllClose(f(), expected) def test_addupdate_on_reshaped_slice(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((4, 3), dtype=jnp.float32)) ref_addupdate(ref.reshape(2, 6), (0, slice(0, 3)), jnp.ones(3, dtype=jnp.float32)) return ref[...] expected = jnp.array( [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=jnp.float32, ) self.assertAllClose(f(), expected) def test_addupdate_on_reshaped_view_slice(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((4, 3), dtype=jnp.float32)) view = ref.reshape(2, 6).at[0, :3] ref_addupdate(view, (), jnp.ones(3, dtype=jnp.float32)) return ref[...] expected = jnp.array( [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=jnp.float32, ) self.assertAllClose(f(), expected) def test_addupdate_on_sliced_ref_reshape(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((4, 6), dtype=jnp.float32)) ref_addupdate(ref.at[0].reshape(2, 3), (), jnp.ones((2, 3), dtype=jnp.float32)) return ref[...] expected = jnp.zeros((4, 6), dtype=jnp.float32).at[0].set(1.0) self.assertAllClose(f(), expected) def test_addupdate_on_reshaped_sliced_reshaped(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((4, 3), dtype=jnp.float32)) view = ref.reshape(2, 6).at[0].reshape(3, 2) ref_addupdate(view, (), jnp.ones((3, 2), dtype=jnp.float32)) return ref[...] expected = jnp.array( [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=jnp.float32, ) self.assertAllClose(f(), expected) def test_addupdate_on_multiple_indexers_raises_error(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((4, 3), dtype=jnp.float32)) ref_addupdate(ref.at[0].at[0], (), jnp.array(1.0, dtype=jnp.float32)) return ref[...] with self.assertRaisesRegex( NotImplementedError, "Multiple indexers are not supported" ): f() def test_addupdate_on_bitcast_raises_error(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((2, 2), dtype=jnp.float32)) ref_addupdate(ref.bitcast(jnp.int32), (), jnp.ones((2, 2), dtype=jnp.int32)) return ref[...] with self.assertRaisesRegex( NotImplementedError, "not supported on bitcast views" ): f() def test_addupdate_on_unsupported_transform_raises_error(self): @jax.jit def f(): ref = jax.new_ref(jnp.zeros((2, 3), dtype=jnp.float32)) view = state_types.TransformedRef( ref, (state_types.TransposeTransform((1, 0)),) ) ref_addupdate(view, (), jnp.ones((3, 2), dtype=jnp.float32)) return ref[...] with self.assertRaisesRegex( NotImplementedError, "Unsupported transform for `addupdate`" ): f() def test_scalar_ref_reshape(self): @jax.jit def f(): ref = jax.new_ref(jnp.array(5.0)) return ref.reshape(1)[...] self.assertEqual(f().shape, (1,)) self.assertAllClose(f(), jnp.array([5.0])) def test_new_ref_host_memory_space_outside_jit(self): ref1 = jax.new_ref( jnp.ones((4,), dtype=jnp.int32), memory_space=core.MemorySpace.Host ) self.assertEqual(core.typeof(ref1).memory_space, core.MemorySpace.Host) self.assertEqual( core.typeof(ref1).inner_aval.memory_space, core.MemorySpace.Device ) self.assertEqual(ref1.sharding.memory_kind, "pinned_host") host_arr = jax.device_put( jnp.arange(4, dtype=jnp.int32), core.MemorySpace.Host ) ref2 = jax.new_ref(host_arr) self.assertEqual(core.typeof(ref2).memory_space, core.MemorySpace.Host) self.assertEqual( core.typeof(ref2).inner_aval.memory_space, core.MemorySpace.Device ) self.assertEqual(ref2.sharding.memory_kind, "pinned_host") np.testing.assert_array_equal( np.asarray(jax.freeze(ref2)), np.arange(4, dtype=np.int32) ) if __name__ == '__main__': absltest.main(testLoader=jtu.JaxTestLoader())