# Copyright 2021 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. import concurrent.futures from functools import partial import unittest from absl.testing import absltest from absl.testing import parameterized import numpy as np import jax from jax import lax from jax.experimental import checkify from jax._src import shard_map from jax.sharding import NamedSharding, PartitionSpec as P from jax._src import array from jax._src import config from jax._src import core from jax._src import test_util as jtu from jax._src.checkify import JaxRuntimeError, FailedCheckError, ErrorEffect, OOBError import jax.numpy as jnp config.parse_flags_with_absl() jtu.request_cpu_devices(8) @jtu.with_config(jax_check_tracer_leaks=True) class CheckifyTransformTests(jtu.JaxTestCase): @jtu.sample_product(jit=[False, True]) @jtu.skip_on_devices("tpu") def test_jit_nan(self, jit): def f(x1, x2): y1 = jnp.sin(x1) y2 = jnp.sin(x2) return y1 + y2 f = jax.jit(f) if jit else f checked_f = checkify.checkify(f, errors=checkify.float_checks) err, _ = checked_f(3., 4.) self.assertIsNone(err.get()) err, _ = checked_f(3., jnp.inf) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: sin") @jtu.sample_product(jit=[False, True]) def test_jit_oob(self, jit): def f(x, i): y = jnp.sin(x) z = y[i] w = jnp.cos(z) return w f = jax.jit(f) if jit else f checked_f = checkify.checkify(f, errors=checkify.index_checks) err, _ = checked_f(jnp.arange(3), 2) self.assertIsNone(err.get()) err, _ = checked_f(jnp.arange(3), 5) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "out-of-bounds indexing") @parameterized.named_parameters( ("get", lambda x: x.get()), ("set", lambda x: x.set(1)), ("add", lambda x: x.add(1)), ("mul", lambda x: x.multiply(1)), ("div", lambda x: x.divide(1)), ("pow", lambda x: x.power(1)), ("min", lambda x: x.min(1)), ("max", lambda x: x.max(1)), ) def test_jit_oob_update(self, update_fn): f = jax.jit(lambda x, i: update_fn(x.at[i])) checked_f = checkify.checkify(f, errors=checkify.index_checks) err, _ = checked_f(jnp.arange(3), 2) self.assertIsNone(err.get()) err, _ = checked_f(jnp.arange(3), 3) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "out-of-bounds indexing") @jtu.sample_product(jit=[False, True]) @jax.numpy_dtype_promotion('standard') def test_jit_div_errors(self, jit): def f(x, y): return x / y f = jax.jit(f) if jit else f checked_f = checkify.checkify(f, errors=checkify.float_checks) err, _ = checked_f(jnp.ones((3,)), jnp.ones((3,))) self.assertIsNone(err.get()) err, _ = checked_f(jnp.ones((3,)), jnp.array([1., 0., 1.])) self.assertIsNotNone(err.get()) err, _ = checked_f(jnp.array([1, jnp.inf, 1]), jnp.array([1, jnp.inf, 1])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: div") @jtu.sample_product(jit=[False, True]) @jtu.skip_on_devices("tpu") def test_jit_multi(self, jit): def f(x, i): y = x[i] z = jnp.cos(y) return z f = jax.jit(f) if jit else f checked_f = checkify.checkify(f, errors=checkify.automatic_checks) # no error err, _ = checked_f(jnp.array([0., jnp.inf, 2.]), 2) self.assertIsNone(err.get()) # oob error err, _ = checked_f(jnp.array([0., 1., 2.]), 5) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "out-of-bounds indexing") # nan error err, _ = checked_f(jnp.array([0., 1., jnp.inf]), 2) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: cos") @parameterized.named_parameters( ("gather", lambda x: x.get()), ("scatter_update", lambda x: x.set(1.)), ("scatter_add", lambda x: x.add(1.)), ("scatter_mul", lambda x: x.multiply(1.)), ("scatter_div", lambda x: x.divide(1.)), ("scatter_pow", lambda x: x.power(1.)), ("scatter_min", lambda x: x.min(1.)), ("scatter_max", lambda x: x.max(1.)), ) def test_numpy_indexing_oobs(self, update_op): def raises_oob(fn, idx, *expected_strs): err, _ = checkify.checkify(jax.jit(fn), errors=checkify.index_checks)(x, idx) error_txt = err.get() self.assertIsNotNone(error_txt) self.assertStartsWith(error_txt, "out-of-bounds indexing") for s in expected_strs: self.assertIn(s, error_txt) x = jnp.ones((2, 3, 7)) axis0_msg = "axis 0 with size 2" axis1_msg = "axis 1 with size 3" axis2_msg = "axis 2 with size 7" single_idx = lambda x, i: update_op(x.at[i]) raises_oob(single_idx, 5, "index 5", axis0_msg) raises_oob(single_idx, -5, "index -3", axis0_msg) raises_oob(single_idx, (0, 100), "index 100", axis1_msg) raises_oob(single_idx, (0, 5, 100), "index 5", axis1_msg) raises_oob(single_idx, (0, 0, 100), "index 100", axis2_msg) raises_oob(single_idx, ((1, 20), (1, 4)), "index 20", axis0_msg) raises_oob(single_idx, ((1, 20), (3, 4)), "index 3", axis1_msg) raises_oob(single_idx, (((1, 1), (1, 20)), 3), "index 3", axis1_msg) raises_oob(single_idx, (((1, 1), (1, 20)), 0), "index 20", axis0_msg) multi_idx = lambda x, i: update_op(x.at[i[0], :, i[1]]) raises_oob(multi_idx, (0, 9), "index 9", axis2_msg) # TODO(lenamartens): numpy reports index -5 here, need to normalize? raises_oob(multi_idx, (-5, 9), "index -3", axis0_msg) raises_oob(multi_idx, (5, -9), "index 5", axis0_msg) raises_oob(multi_idx, ((0, 9), 0), "index 9", axis0_msg) def test_dynamic_slice_oobs(self): def raises_oob(fn, x, idx, *expected_strs): err, _ = checkify.checkify(jax.jit(fn), errors=checkify.index_checks)(x, idx) error_txt = err.get() self.assertIsNotNone(error_txt) self.assertStartsWith(error_txt, "out-of-bounds indexing") for s in expected_strs: self.assertIn(s, error_txt) x = jnp.ones((2, 3, 7)) raises_oob(partial(lax.dynamic_slice, slice_sizes=(1, 1, 1)), x, (2, 0, 0), 'index 2') raises_oob(partial(lax.dynamic_slice, slice_sizes=(1, 1, 1)), x, (-3, 0, 0), 'index -1') raises_oob(partial(lax.dynamic_slice, slice_sizes=(1, 1, 1)), x, (0, 3, 0), 'index 3') raises_oob(partial(lax.dynamic_slice, slice_sizes=(1, 1, 1)), x, (0, -5, 0), 'index -2') raises_oob(partial(lax.dynamic_slice, slice_sizes=(1, 1, 1)), x, (0, 1, 8), 'index 8') raises_oob(partial(lax.dynamic_slice, slice_sizes=(1, 1, 1)), x, (0, 1, -10), 'index -3') def test_dynamic_update_slice_oobs(self): def raises_oob(fn, x, y, idx, *expected_strs): err, _ = checkify.checkify(jax.jit(fn), errors=checkify.index_checks)(x, y, idx) error_txt = err.get() self.assertIsNotNone(error_txt) self.assertStartsWith(error_txt, "out-of-bounds indexing") for s in expected_strs: self.assertIn(s, error_txt) x = jnp.ones((2, 3, 7)) y = jnp.zeros((1, 1, 1)) raises_oob(lax.dynamic_update_slice, x, y, (2, 0, 0), 'index 2') raises_oob(lax.dynamic_update_slice, x, y, (-3, 0, 0), 'index -1') raises_oob(lax.dynamic_update_slice, x, y, (0, 3, 0), 'index 3') raises_oob(lax.dynamic_update_slice, x, y, (0, -5, 0), 'index -2') raises_oob(lax.dynamic_update_slice, x, y, (0, 1, 8), 'index 8') raises_oob(lax.dynamic_update_slice, x, y, (0, 1, -10), 'index -3') @jtu.sample_product(jit=[False, True]) def test_jit_ordering(self, jit): def f(x, i): y = x[i] z = jnp.sin(x) return y * z f = jax.jit(f) if jit else f checked_f = checkify.checkify(f, errors=checkify.automatic_checks) # both oob and nan error, but oob happens first err, _ = checked_f(jnp.array([0., 1., jnp.inf]), 5) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "out-of-bounds indexing") def test_pmap_basic(self): if len(jax.devices()) < 2: raise unittest.SkipTest("requires at least 2 devices") @jax.pmap def f(x): y1 = jnp.sin(1./x) y2 = jnp.sin(x) return y1 + y2 checked_f = checkify.checkify(f, errors=checkify.nan_checks) xs = jnp.array([1., 2.]) err, _ = checked_f(xs) self.assertIsNone(err.get()) xs = jnp.array([3., 0.]) err, _ = checked_f(xs) self.assertIsNotNone(err.get()) self.assertIn("nan generated by primitive: sin", err.get()) def test_pmap_collectives(self): if len(jax.devices()) < 4: raise unittest.SkipTest("requires at least 2 devices") @partial(jax.pmap, axis_name="i") def f(x1): return jax.lax.all_gather(x1, axis_name="i") checked_f = checkify.checkify(f, errors=checkify.float_checks) xs = jnp.array([0., 2., 3., 6.]) err, _ = checked_f(xs) self.assertIsNone(err.get()) @jtu.skip_on_devices("tpu") def test_cond_basic(self): @jax.jit def f(x): def true_fun(x): return jnp.sin(x) def false_fun(x): checkify.check(x > -1, "oh no") return x / 0. return lax.cond(x > 0, true_fun, false_fun, x) checked_f = checkify.checkify(f, errors=checkify.all_checks) err, _ = checked_f(3.) self.assertIsNone(err.get()) err, _ = checked_f(jnp.inf) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: sin") err, _ = checked_f(-jnp.inf) self.assertStartsWith(err.get(), "oh no") err, _ = checked_f(0.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") def test_cond_different_payloads(self): @jax.jit def f(x): def true_fun(x): checkify.check(~x, "{one}", one=x) def false_fun(x): checkify.check(x, "{one} and {two}", one=x, two=x) return lax.cond(x, true_fun, false_fun, x) checked_f = checkify.checkify(f) err, _ = checked_f(True) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "True") err, _ = checked_f(False) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "False and False") def test_cond_nd_payloads(self): @jax.jit def f(x): def true_fun(x): checkify.check(jnp.all(x > 0), "{one}", one=x) def false_fun(x): checkify.check(jnp.all(x < 0), "{one} and {two}", one=x, two=x) return lax.cond(jnp.all(x < 0), true_fun, false_fun, x) checked_f = checkify.checkify(f) err, _ = checked_f(jnp.arange(0, 4)) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "[0 1 2 3] and [0 1 2 3]") err, _ = checked_f(jnp.arange(-4, -1)) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "[-4 -3 -2]") @jtu.skip_on_devices("tpu") def test_scan_map(self): def scan_body(_, x): return None, jnp.sin(x) @jax.jit def f(xs): return lax.scan(scan_body, None, xs) checked_f = checkify.checkify(f, errors=checkify.float_checks) xs = jnp.array([0., 2.]) err, (_, ch_outs) = checked_f(xs) _, outs = f(xs) self.assertIsNone(err.get()) self.assertArraysEqual(ch_outs, outs) xs = jnp.array([3., jnp.inf]) err, (_, ch_outs) = checked_f(xs) _, outs = f(xs) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: sin") self.assertArraysEqual(ch_outs, outs) @jtu.skip_on_devices("tpu") def test_scan_carry(self): def scan_body(carry, x): carry = carry-1. possible_nan = jnp.sin(1./carry) return carry, x+possible_nan @jax.jit def f(carry, xs): return lax.scan(scan_body, carry, xs) checked_f = checkify.checkify(f, errors=checkify.float_checks) carry, xs = 3., jnp.ones((2,)) err, (ch_out_carry, ch_outs) = checked_f(carry, xs) out_carry, outs = f(carry, xs) self.assertIsNone(err.get()) self.assertArraysEqual(ch_outs, outs) self.assertArraysEqual(ch_out_carry, out_carry) # error happens on first iteration carry, xs = 1., jnp.ones((2,)) err, (ch_out_carry, ch_outs) = checked_f(carry, xs) out_carry, outs = f(carry, xs) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") self.assertArraysEqual(ch_outs, outs) self.assertArraysEqual(ch_out_carry, out_carry) # error happens on second iteration carry, xs = 2., jnp.ones((4,)) err, (ch_out_carry, ch_outs) = checked_f(carry, xs) out_carry, outs = f(carry, xs) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") self.assertArraysEqual(ch_outs, outs) self.assertArraysEqual(ch_out_carry, out_carry) @jtu.skip_on_devices("tpu") def test_while_loop_body_error(self): def while_cond(val): i, _ = val return i < 2 def while_body(val): i, x = val possible_nan = jnp.sin(1./i) return i+1., x+possible_nan @jax.jit def f(init_val): return lax.while_loop(while_cond, while_body, (init_val, 0.)) checked_f = checkify.checkify(f, errors=checkify.float_checks) init_val = 1. err, ch_out = checked_f(init_val) out = f(init_val) self.assertIsNone(err.get()) self.assertArraysEqual(ch_out, out) init_val = 0. err, ch_out = checked_f(init_val) out = f(init_val) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") self.assertArraysEqual(ch_out, out) @jtu.skip_on_devices("tpu") def test_while_loop_cond_error(self): def while_cond(val): _ = jnp.sin(1./val) return 0. * _ + val < 2. def while_body(val): return val+1. @jax.jit def f(init_val): return lax.while_loop(while_cond, while_body, init_val) checked_f = checkify.checkify(f, errors=checkify.float_checks) init_val = 1. err, ch_out = checked_f(init_val) out = f(init_val) self.assertIsNone(err.get()) self.assertArraysEqual(ch_out, out) init_val = 0. err, ch_out = checked_f(init_val) out = f(init_val) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") self.assertArraysEqual(ch_out, out) @jtu.skip_on_devices("tpu") def test_while_loop_cond_error_and_false(self): # Tests if an error is generated when cond returns False. def while_cond(val): possible_nan = jnp.sin(1./val) return jnp.logical_not(jnp.isnan(possible_nan)) @jax.jit def f(init_val): return lax.while_loop(while_cond, lambda val: val-1, init_val) checked_f = checkify.checkify(f, errors=checkify.float_checks) # error on first cond init_val = 0. err, _ = checked_f(init_val) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") # error on second cond init_val = 1. err, _ = checked_f(init_val) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") def test_checify_donation_no_forwarding(self): mesh = jtu.create_mesh((2,), ('x',)) @checkify.checkify @jax.jit(donate_argnums=(0,)) def f(x: jax.Array) -> jax.Array: checkify.check(jnp.all(x > 0), "a") return x x = jax.device_put(jnp.zeros(64, dtype="int32"), NamedSharding(mesh, P())) err, y = f(x) err, z = f(y) # doesn't crash @jtu.with_explicit_mesh((2,), "x") def test_auto_axes_explicit_top_level(self, mesh): x = jnp.ones((4,), out_sharding=jax.P("x")) @jax.jit @checkify.checkify def my_func(x): return x.at[:].set(2.0, out_sharding=jax.P("x")) my_func(x) @jtu.skip_on_devices("tpu") def test_while_loop_body_and_cond_error(self): def while_cond(val): i, cond_val, _ = val j = jnp.sin(cond_val) return i + (0. * j) < 2 # don't let the sin value be dead code def while_body(val): i, cond_val, body_val = val possible_nan = jnp.cos(body_val) return i+1., cond_val, possible_nan @jax.jit def f(cond_val, body_val): return lax.while_loop(while_cond, while_body, (0., cond_val, body_val)) checked_f = checkify.checkify(f, errors=checkify.float_checks) cond_val = jnp.inf body_val = 1. err, _ = checked_f(cond_val, body_val) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: sin") cond_val = 1. body_val = jnp.inf err, _ = checked_f(cond_val, body_val) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: cos") cond_val = jnp.inf body_val = jnp.inf err, _ = checked_f(cond_val, body_val) self.assertIsNotNone(err.get()) # first error which occurs is in cond self.assertStartsWith(err.get(), "nan generated by primitive: sin") def test_checkify_jit(self): def f(x): # unary func return x / x def g(x, y): # binary func return x / y devices = jax.local_devices()[:8] # Taking up to 8 devices mesh = jax.sharding.Mesh(np.array(devices), ["dev"]) ps = NamedSharding(mesh, jax.sharding.PartitionSpec("dev")) inp = np.arange(8) x = array.make_array_from_callback(inp.shape, ps, lambda idx: inp[idx]) f = jax.jit(f, in_shardings=ps, out_shardings=ps) f = checkify.checkify(f, errors=checkify.float_checks) g = jax.jit(g, in_shardings=ps, out_shardings=ps) g = checkify.checkify(g, errors=checkify.float_checks) with jax.set_mesh(mesh): u_err, _ = f(x) b_err, _ = g(x, x) self.assertIsNotNone(u_err.get()) self.assertStartsWith(u_err.get(), "division by zero") self.assertIsNotNone(b_err.get()) self.assertStartsWith(b_err.get(), "division by zero") @parameterized.parameters(True, False) def test_shard_map(self, check_vma): def f(x): # unary func return jax.lax.axis_index("dev") * x / x def g(x, y): # binary func return jax.lax.axis_index("dev") * x / y devices = jax.local_devices()[:8] # Taking up to 8 devices mesh = jax.sharding.Mesh(np.array(devices), ["dev"]) pspec = jax.sharding.PartitionSpec("dev") ps = NamedSharding(mesh, pspec) inp = np.tile(np.arange(4, dtype=np.int32), 2) x = array.make_array_from_callback(inp.shape, ps, lambda idx: inp[idx]) f = shard_map.shard_map( f, mesh=mesh, in_specs=pspec, out_specs=pspec, check_vma=check_vma ) f = jax.jit(f, in_shardings=ps, out_shardings=ps) f = checkify.checkify(f, errors=checkify.float_checks) g = shard_map.shard_map( g, mesh=mesh, in_specs=(pspec, pspec), out_specs=pspec, check_vma=check_vma ) g = jax.jit(g, in_shardings=(ps, ps), out_shardings=ps) g = checkify.checkify(g, errors=checkify.float_checks) u_err, _ = f(x) b_err, _ = g(x, x) divbyzero = "division by zero" expected_err = f"at mapped index 0: {divbyzero}" if (next_device_with_zero := len(devices) // 2) != 0: expected_err += f"\nat mapped index {next_device_with_zero}: {divbyzero}" self.assertIsNotNone(u_err.get()) self.assertEqual(u_err.get(), expected_err) self.assertIsNotNone(b_err.get()) self.assertEqual(b_err.get(), expected_err) def test_empty_enabled_errors(self): def multi_errors(x): x = x/0 # DIV x = jnp.sin(x) # NAN x = x[500] # OOB checkify.check(x < 0, "must be negative!") # ASSERT return x x = jnp.ones((2,)) err, _ = checkify.checkify(multi_errors, errors=set())(x) self.assertIsNone(err.get()) @parameterized.named_parameters( ("assert", checkify.user_checks, "must be negative!"), ("div", checkify.div_checks, "division by zero"), ("nan", checkify.nan_checks, "nan generated"), ("oob", checkify.index_checks, "out-of-bounds indexing"), ("automatic_checks", checkify.automatic_checks, "division by zero"), ) @jtu.skip_on_devices("tpu") def test_enabled_errors(self, error_set, expected_error): def multi_errors(x): checkify.check(jnp.all(x < 0), "must be negative!") # ASSERT x = x/0 # DIV x = jnp.sin(x) # NAN x = x[500] # OOB return x x = jnp.ones((2,)) err, _ = checkify.checkify(multi_errors, errors=error_set)(x) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), expected_error) @jtu.skip_on_devices("tpu") def test_post_process_call(self): @partial(checkify.checkify, errors=checkify.float_checks) def g(x): @jax.jit def f(y): return jnp.sin(x * y) return f(jnp.inf) err, _ = g(2.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: sin") @jtu.skip_on_devices("tpu") def test_post_process_map(self): @partial(checkify.checkify, errors=checkify.float_checks) def g(x): @jax.pmap def f(y): return jnp.sin(x * y), jnp.cos(x * y) return f(jnp.array([jnp.inf]))[0] err, _ = g(2.) self.assertIsNotNone(err.get()) self.assertIn("nan generated by primitive: sin", err.get()) @jtu.skip_on_devices("tpu") def test_custom_jvp(self): @jax.custom_jvp def sin(x): return jnp.sin(x) @sin.defjvp def sin_jvp(primals, tangents): (x,), (xdot,) = primals, tangents return sin(x), jnp.cos(x) * xdot f = checkify.checkify(sin, errors=checkify.float_checks) err, y = f(3.) self.assertIsNone(err.get()) err, y = f(jnp.inf) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), 'nan generated by primitive: sin') # When we hit the custom jvp rule with jvp-of-checkify, no checks are added. (err, y), (errdot, ydot) = jax.jvp(f, (3.,), (1.,)) # doesn't crash self.assertIsNone(err.get()) # no error if config.config.jax_custom_jvp3: self.assertNotEmpty(err._metadata) else: self.assertEmpty(err._metadata) # and no checks were added! self.assertEmpty(errdot._metadata) y_expected, ydot_expected = jax.jvp(jnp.sin, (3.,), (1.,)) self.assertAllClose(y, y_expected) self.assertAllClose(ydot, ydot_expected) # Grad-of-checkify doesn't crash either. x_bar = jax.grad(lambda x: f(x)[1])(3.) self.assertAllClose(x_bar, jnp.cos(3.)) # Checkify-of-jvp adds checks (unlike jvp-of-checkify above). g = checkify.checkify(lambda x, xdot: jax.jvp(sin, (x,), (xdot,)), errors=checkify.float_checks) err, (y, ydot) = g(3., 1.) # doesn't crash self.assertIsNone(err.get()) # no error self.assertNotEmpty(err._metadata) # but checks were added! self.assertAllClose(y, jnp.sin(3.)) self.assertAllClose(ydot, jnp.cos(3.)) err, _ = g(jnp.inf, 1.) self.assertIsNotNone(err.get()) # yes error self.assertStartsWith(err.get(), 'nan generated by primitive: sin') @jtu.skip_on_devices("tpu") def test_custom_jvp_vmap(self): @jax.custom_jvp def log(x): return jnp.log(x) log.defjvps(lambda t, ans, x: t / x) f = checkify.checkify(jax.vmap(log), errors=checkify.float_checks) err, _ = f(jnp.array([1., -1.])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), 'nan generated by primitive: log') def test_custom_vjp(self): @jax.custom_vjp def sin(x): return jnp.sin(x) def sin_fwd(x): return jnp.sin(x), 2. * x def sin_bwd(x2, g): return jnp.cos(x2 / 2.) * g, sin.defvjp(sin_fwd, sin_bwd) f = checkify.checkify(sin, errors=checkify.float_checks) # no differentiation, no error err, y = f(3.) self.assertIsNone(err.get()) # no differentiation, yes error err, y = f(jnp.inf) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), 'nan generated by primitive: sin') # When we hit the custom vjp rule with vjp-of-checkify, no checks are added. (err, y), f_vjp = jax.vjp(f, 3.) self.assertIsNone(err.get()) # no error if config.config.jax_custom_vjp3: self.assertNotEmpty(err._metadata) else: self.assertEmpty(err._metadata) # and no checks were added! # Checkify-of-vjp adds checks (unlike vjp-of-checkify above). err, y = checkify.checkify(jax.grad(sin), errors=checkify.float_checks)(3.) self.assertIsNone(err.get()) # no error self.assertNotEmpty(err._metadata) # but checks were added! err, y = checkify.checkify(jax.grad(sin), errors=checkify.float_checks)(jnp.inf) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive: sin") @parameterized.parameters('jvp', 'vjp') @config.custom_jvp3(True) @config.custom_vjp3(True) def test_hijax_custom_derivative_primal_error(self, kind): if kind == 'jvp': f = jax.custom_jvp(jnp.log) f.defjvp(lambda p, t: (f(p[0]), t[0] / p[0])) else: f = jax.custom_vjp(jnp.log) f.defvjp(lambda x: (f(x), x), lambda x, g: (g / x,)) checked = checkify.checkify(f, errors=checkify.float_checks) err, _ = checked(-1.) self.assertStartsWith(err.get(), 'nan generated by primitive: log') err, _ = checked(1.) self.assertIsNone(err.get()) self.assertAllClose(jax.grad(lambda x: checked(x)[1])(2.), 0.5) @jtu.skip_on_devices("tpu") def test_custom_vjp_explicit_check(self): @jax.custom_vjp def f(x): checkify.check(x > 0, "check in f") return x * 2. def f_fwd(x): checkify.check(x > -10, "check in fwd") return f(x), (x,) def f_bwd(res, g): checkify.check(g > 0, "check in bwd") return (g * 2.,) f.defvjp(f_fwd, f_bwd) # 1. Primal check under checkify: # Under custom_vjp3, ErrorEffect is disallowed in the primal function, # raising NotImplementedError at trace time. checked_f = checkify.checkify(f) if not config.config.jax_custom_vjp3: err, _ = jax.jit(checked_f)(-1.) self.assertIn("check in f", err.get()) else: with self.assertRaisesRegex( NotImplementedError, "checkify effects .* are not supported in custom_vjp-decorated primal" " functions", ): checked_f(-1.0) # 2. Checks in fwd and bwd are supported under both classic and custom_vjp3: @jax.custom_vjp def g(x): return x * 2.0 def g_fwd(x): checkify.check(x > -10, "check in fwd") return x * 2.0, (x,) def g_bwd(res, cotangent): checkify.check(cotangent > 0, "check in bwd") return (cotangent * 2.0,) g.defvjp(g_fwd, g_bwd) checked_grad = checkify.checkify(jax.grad(g)) err, _ = jax.jit(checked_grad)( 1.0 ) # all checks pass (x=1 > -10, cotangent=1 > 0) self.assertIsNone(err.get()) err, _ = jax.jit(checked_grad)(-20.) # fails "check in fwd" (x=-20 < -10) self.assertIn("check in fwd", err.get()) _, vjp_fn = jax.vjp(g, 1.0) err, _ = jax.jit(checkify.checkify(vjp_fn))( -1.0 ) # cotangent=-1 fails "check in bwd" self.assertIn("check in bwd", err.get()) # 3. If a primal check is placed outside the custom_vjp decorator, it works: def g_safe(x): checkify.check(x > 0, "check outside custom_vjp") return g(x) err, _ = jax.jit(checkify.checkify(g_safe))(-1.0) self.assertIn("check outside custom_vjp", err.get()) def test_scan_consts(self): def f(xs): def scan_body(carry, _): # closes oves xs return carry+1, xs[carry] return lax.scan(scan_body, 1, xs) checked_f = checkify.checkify(f, errors=checkify.index_checks) err, _ = checked_f(jnp.ones((7,))) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "out-of-bounds indexing") def test_scan_consts2(self): def f(xs): def scan_body(carry, _): # add more consts! _ = xs[carry], xs[carry], jnp.sin(np.arange(11.)) return carry+1, xs[carry] return lax.scan(scan_body, 1, xs)[1] checked_f = checkify.checkify(f, errors=checkify.index_checks) err, _ = checked_f(jnp.ones((7, 3))) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "out-of-bounds indexing") def test_while_consts(self): def f(xs): def while_cond(carry): i, _ = carry _ = xs[i], jnp.sin(np.arange(11.)) return i > -1 def while_body(carry): i, _ = carry x = xs[i] return i - 1, x/i return lax.while_loop(while_cond, while_body, (0, jnp.zeros_like(xs[0]))) checked_f = checkify.checkify(f, errors=checkify.float_checks) err, _ = checked_f(jnp.ones((7, 3))) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "division by zero") def test_multiple_payloads(self): def f(x): a = x[5] b = x[6] return a + b err, _ = checkify.checkify(f, errors=checkify.index_checks)(jnp.ones((2,))) self.assertIsNotNone(err.get()) self.assertIn("index 5", err.get()) def test_nd_payloads(self): cf = checkify.checkify(lambda x, i: x[i], errors=checkify.index_checks) errs, _ = jax.vmap(cf)(jnp.ones((3, 2)), jnp.array([5, 0, 100])) self.assertIsNotNone(errs.get()) self.assertIn("index 5", errs.get()) self.assertIn("index 100", errs.get()) def test_mapped_error_one_payload(self): def f(x, i): x = x[i] return x/0 cf = checkify.checkify(f, errors=checkify.automatic_checks) errs, _ = jax.vmap(cf)(jnp.ones((2, 1)), jnp.array([0, 100])) self.assertIsNotNone(errs.get()) self.assertIn("division by zero", errs.get()) self.assertIn("index 100", errs.get()) @jax.legacy_prng_key('allow') def test_checking_key_split_with_nan_check(self): cf = checkify.checkify( lambda k: jax.random.permutation(k, jnp.array([0, 1, 2])), errors=checkify.float_checks) cf(jax.random.PRNGKey(123)) # does not crash. def test_pmap_one_device(self): @jax.pmap def f(x, y): return x/y cf = checkify.checkify(f, errors=checkify.automatic_checks) errs, _ = cf(jnp.ones((1,)), jnp.zeros((1,))) self.assertIsNotNone(errs.get()) self.assertIn("division by zero", errs.get()) def test_psum_nan_check(self): @partial(jax.vmap, axis_name="i") def f(x, y): return lax.psum((x, y), axis_name="i") cf = checkify.checkify(f, errors=checkify.nan_checks) err, _ = cf(jnp.array([-jnp.inf, 0, jnp.inf]), jnp.ones((3, 2))) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "nan generated by primitive") def test_different_payload_effects(self): def f(x, y): x = x[y] checkify.check(jnp.all(x > 0), "{x}", x=x) return x f = checkify.checkify(f, errors=checkify.all_checks) err, _ = jax.vmap(f)(jnp.ones((2, 3))*-1, jnp.array([0, 5])) self.assertIsNotNone(err.get()) def test_effects_total_ordering(self): sds0 = jax.ShapeDtypeStruct((2,), jnp.float32) sds1 = jax.ShapeDtypeStruct((2,), jnp.int32) sds2 = jax.ShapeDtypeStruct((3,), jnp.int32) self.assertTotallyOrdered( [ErrorEffect(FailedCheckError, (sds0,))], [ErrorEffect(FailedCheckError, (sds0, sds0))], [ErrorEffect(FailedCheckError, (sds1,))], [ErrorEffect(FailedCheckError, (sds1, sds0))], [ErrorEffect(FailedCheckError, (sds2,))], [ErrorEffect(OOBError, (sds0,))], [ErrorEffect(OOBError, (sds0, sds0))], ) def test_scan_xs_mapped_correctly(self): def f(_, x): return None, jnp.reshape(x, (2, 2)) @jax.jit def g(x): return jax.lax.scan(f, None, x) checked_f = checkify.checkify(g) checked_f = jax.jit(checked_f) err, _ = checked_f(jnp.ones((2, 4))) self.assertIsNone(err.get()) def test_retracing(self): f = checkify.checkify(jax.jit(lambda x: jnp.sin(x) ** 2)) _ = f(3.) with jtu.count_jit_and_pmap_lowerings() as count: _ = f(3.) self.assertEqual(count(), 0) def test_goodfellow_custom_jvp(self): def h(fext): checkify.check(True, "") return jax.nn.relu(fext) h = checkify.checkify(h) def h_out(fext): _, out = h(fext) return out h_grad = jax.grad(h_out) h_grad(0.) # doesn't crash def test_goodfellow_custom_vjp(self): @jax.custom_vjp def sin(x): return jnp.sin(x) def sin_fwd(x): return jnp.sin(x), 2. * x def sin_bwd(x2, g): return jnp.cos(x2 / 2.) * g, sin.defvjp(sin_fwd, sin_bwd) def h(fext): checkify.check(True, "") return sin(fext) h = checkify.checkify(h) def h_out(fext): _, out = h(fext) return out h_grad = jax.grad(h_out) h_grad(0.) # doesn't crash def test_closed_call(self): # lots of golfing went into this test y = jnp.array([3.14]) summify = lambda f: lambda x: f(x).sum() f = checkify.checkify(jax.grad(summify(jax.remat( partial(partial, jax.lax.map)(lambda x: jnp.sin(x * y)))))) f(jnp.array([3.])) # don't crash def test_while_loop_leaks(self): def f(x): n = jnp.minimum(1, 2) return jax.lax.while_loop(lambda i: i < n, lambda i: i + 1, x) jax.jit(checkify.checkify(f))(0) # Does not crash bc of leaked tracer. @parameterized.parameters(True, False) def test_remat(self, jit): # basic test from https://github.com/jax-ml/jax/issues/23867 def fn(x: jax.Array): checkify.check(jnp.all(x > 0), "x must be positive") return x + 1 fn = jax.remat(fn) if jit: fn = jax.jit(fn) fn = checkify.checkify(fn) err, y = fn(jnp.array([1, 2, 3])) self.assertIsNone(err.get()) self.assertAllClose(y, jnp.array([2, 3, 4]), check_dtypes=False) err, _ = fn(jnp.array([0, 2, 3])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "x must be positive") def test_checkify_custom_jvp(self): @jax.custom_jvp def f(x): return x * 2.0 @f.defjvp def f_jvp(primals, tangents): x, = primals dx, = tangents return f(x), dx * 2.0 def body(x): checkify.check(x > 0, "x must be positive") return f(x) checked_f = jax.jit(checkify.checkify(body)) err, res = checked_f(2.0) self.assertIsNone(err.get()) self.assertAllClose(res, 4.0) @jtu.with_config(jax_check_tracer_leaks=True) class AssertPrimitiveTests(jtu.JaxTestCase): def test_assert_primitive_impl(self): def f(): checkify.check(False, "hi") with self.assertRaisesRegex(JaxRuntimeError, "hi"): f() def test_assert_primitive_lowering(self): @jax.jit def f(): checkify.check(False, "hi") with self.assertRaisesRegex(ValueError, "Cannot abstractly evaluate"): f() def test_assert_primitive_jaxpr_effects(self): def f(x): checkify.check(False, "hi: {}", x) jaxpr = jax.make_jaxpr(f)(jnp.ones(4, jnp.int32)) self.assertSetEqual(jaxpr.effects, {ErrorEffect(FailedCheckError, ( jax.ShapeDtypeStruct((4,), jnp.int32),))}) def g(x, y): checkify.check(False, "hi: {} {}", x, y) self.assertSetEqual( jax.make_jaxpr(g)( jnp.ones(4, jnp.int32), jnp.ones(2, jnp.float32)).effects, {ErrorEffect(FailedCheckError, ( jax.ShapeDtypeStruct((4,), jnp.int32), jax.ShapeDtypeStruct((2,), jnp.float32)))}) def test_assert_primitive_eval_shape(self): # The check is abstractly evaluated but not lowered. def f(): checkify.check(False, "hi") jax.eval_shape(f) # does not crash. def test_assert_discharging(self): @checkify.checkify def f(x): checkify.check(x > 0, "must be positive!") return jnp.log(x) err, _ = f(1.) self.assertIsNone(err.get()) err, _ = f(0.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "must be positive") f = jax.jit(f) err, _ = f(1.) self.assertIsNone(err.get()) err, _ = f(0.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "must be positive") def test_assert_discharging_no_data_dependence(self): @jax.jit def g(x): @checkify.checkify def f(): # Note that x is not an argument to the checkified function. checkify.check(x > 0, "must be positive!") return jnp.log(x) return f() err, _ = g(1.) self.assertIsNone(err.get()) err, _ = g(0.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "must be positive") def test_assert_discharging_scan(self): def body(carry, x): checkify.check(jnp.all(x > 0), "must be positive") return carry, x def f(x): return jax.lax.scan(body, (None,), x) err, _ = checkify.checkify(f)(jnp.array([-1])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "must be positive") err, _ = checkify.checkify(f)(jnp.array([1, 0, -1])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "must be positive") def test_assert_discharging_while_loop(self): def while_cond(val): i, _ = val checkify.check(i < 0, "i must be negative") return i < 2 def while_body(val): i, x = val checkify.check(x < 0, "x must be negative") return i+1., x+1 @jax.jit def f(init_i, init_val): return lax.while_loop(while_cond, while_body, (init_i, init_val)) checked_f = checkify.checkify(f) err, _ = checked_f(0, 1) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "i must be negative") err, _ = checked_f(-1, 0) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "x must be negative") def test_assert_discharging_cond(self): def true_branch(x): checkify.check(jnp.all(x != 0.), "x cannot be 0") return 1/x def false_branch(x): checkify.check(jnp.all(x >= 0), "x must be positive") return x*2 @jax.jit def f(pred, x): return lax.cond(pred, true_branch, false_branch, x) checked_f = checkify.checkify(f) err, _ = checked_f(True, 0.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "x cannot be 0") err, _ = checked_f(False, 0.) self.assertIsNone(err.get()) err, _ = checked_f(False, -1.) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "x must be positive") err, _ = checked_f(True, -1.) self.assertIsNone(err.get()) def test_assert_batching_rule(self): @jax.vmap def f(x): checkify.check(jnp.sum(x) == 1., "x must sum to one.") return x no_failures = jnp.array([[0.5, 0.5], [1., 0.]]) one_batch_fails = jnp.array([[0.5, 0.5], [1, 1]]) mult_batch_fail = jnp.array([[0.5, 0.5], [1, 1], [2, 2]]) f(no_failures) with self.assertRaisesRegex(JaxRuntimeError, "x must sum to one."): f(one_batch_fails) with self.assertRaisesRegex(JaxRuntimeError, "x must sum to one."): f(mult_batch_fail) checked_f = checkify.checkify(f) err, _ = checked_f(no_failures) self.assertIsNone(err.get()) err, _ = checked_f(one_batch_fails) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "x must sum to one") err, _ = checked_f(mult_batch_fail) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "x must sum to one") def test_check_error(self): def g(): checkify.check(False, "hi") def f(): err, _ = checkify.checkify(g)() checkify.check_error(err) with self.assertRaisesRegex(JaxRuntimeError, "hi"): f() f = checkify.checkify(f) err, none = f() self.assertIsNone(none) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "hi") def test_check_error_scanned(self): def body(carry, x): checkify.check(jnp.all(x > 0), "should be positive") return carry, x def checked_body(carry, x): err, (carry, x) = checkify.checkify(body)(carry, x) return carry, (x, err) def f(x): _, (xs, errs) = jax.lax.scan(checked_body, (None,), x) checkify.check_error(errs) return xs err, _ = checkify.checkify(f)(jnp.array([-1])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "should be positive") err, _ = checkify.checkify(f)(jnp.array([1, 0, -1])) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "should be positive") def test_discharge_recharge(self): def ejit(f): f = checkify.checkify(f) f = jax.jit(f) def jitted_f(*args): err, out = f(*args) checkify.check_error(err) return out return jitted_f @ejit def f(pred): assert python_should_be_running checkify.check(pred, "foo") python_should_be_running = True f(True) python_should_be_running = False f(True) with self.assertRaisesRegex(JaxRuntimeError, "foo"): f(False) def test_cond_of_named_call(self): def g(x): branch = jax.named_call(lambda x: x) out = jax.lax.cond(True, branch, branch, x) return out checkify.checkify(g)(0.) # does not crash def test_grad(self): @jax.grad def f(x): checkify.check(jnp.all(x > 0), "should be positive!") return x f = checkify.checkify(f) err, _ = f(1.) self.assertIsNone(err.get()) err, _ = f(0.) self.assertIsNotNone(err.get()) self.assertIn("should be positive", err.get()) def test_checkify_of_vmap_of_while_errors(self): @jax.vmap def fun(n, v): def while_cond(s): counter, value = s checkify.check(value < 6, "value needs to be less than 6!") return counter > 0 def while_body(s): counter, value = s checkify.check(value >= 0, "value needs to be positive!") return counter/value, value - 1. _, result = jax.lax.while_loop(while_cond, while_body, (n, v)) return result checked_f = checkify.checkify(fun, errors=checkify.all_checks) with self.assertRaisesRegex(ValueError, "checkify-of-vmap-of-while"): checked_f(jnp.asarray([1., 2., 3.]), jnp.asarray([5., 2., 4.])) # TODO(lenamartens): re-enable assertions below. # self.assertIsNotNone(err.get()) # self.assertStartsWith(err.get(), "division by zero") # err, _ = checked_f(jnp.asarray([1., 2., 3.]), jnp.asarray([5., 2., -4.])) # self.assertIsNotNone(err.get()) # self.assertStartsWith(err.get(), "value needs to be positive") # err, _ = checked_f(jnp.asarray([1., 2., 3.]), jnp.asarray([6., 2., -4.])) # self.assertIsNotNone(err.get()) # self.assertStartsWith(err.get(), "value needs to be less than 6") def test_checkify_of_vmap_of_while_masked_errors(self): def cond(x): return x < 5 def body(x): # This will only trigger in the masked portion of the batched while. checkify.check(x < 5, "should never happen") return x + 1 @jax.vmap def fun(x): return lax.while_loop(cond, body, x) checked_f = checkify.checkify(fun) with self.assertRaisesRegex(ValueError, "checkify-of-vmap-of-while"): checked_f(jnp.arange(5)) # TODO(lenamartens): re-enable assertions below. # self.assertIsNone(err.get()) def test_vmap_of_checkify_of_while(self): # vmap-of-checkify-of-while is a supported composition: checkify injects a # dce_sink into the while body and the outer vmap must batch it. Regression # test for "foreach() argument 2 is longer than argument 1". def fun(v): def while_cond(s): counter, value = s return value < 6. def while_body(s): counter, value = s checkify.check(value >= 0, "value needs to be positive!") return counter + 1, value + 1. counter, _ = jax.lax.while_loop(while_cond, while_body, (jnp.int32(0), v)) return counter checked_f = checkify.checkify(fun, errors=checkify.all_checks) err, counts = jax.jit(jax.vmap(checked_f))(jnp.asarray([1., 2., 3.])) self.assertIsNone(err.get()) self.assertArraysEqual(counts, jnp.asarray([5, 4, 3], dtype=jnp.int32)) def test_assert_cond_no_data_dependence(self): def true_fun(): return checkify.check(False, "hi!") def false_fun(): return checkify.check(False, "bye!") def f(): return jax.lax.cond(True, true_fun, false_fun) f = checkify.checkify(f) err, _ = f() self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "hi!") def test_assert_switch_no_data_dependence(self): def branch(): checkify.check(False, "hi!") def f(): return lax.switch(0, [branch]*3) checked_f = checkify.checkify(f) err, _ = checked_f() self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "hi!") def test_debug_check_noop(self): def f(x): checkify.debug_check(jnp.all(x != x), "{x} cannot be {x}", x=x) return x x = jnp.ones(()) f(x) # no error. jax.jit(f)(x) # no error. jax.vmap(f)(jnp.ones((2,))) # no error. jax.grad(f)(x) # no error. @parameterized.named_parameters(("with_jit", True), ("without_jit", False)) def test_debug_check_nonscalar_pred(self, with_jit): def f(x): checkify.debug_check(x != x, "{x} cannot be {x}", x=x) return x checked_f = checkify.checkify(f) if with_jit: checked_f = jax.jit(checked_f) with self.assertRaisesRegex(TypeError, "debug_check takes a scalar pred"): checked_f(jnp.ones((5,))) @parameterized.named_parameters(("with_jit", True), ("without_jit", False)) def test_debug_check(self, with_jit): def f(x): checkify.debug_check(jnp.all(x != x), "{x} cannot be {x}", x=x) return x checked_f = checkify.checkify(f) if with_jit: checked_f = jax.jit(checked_f) err, _ = checked_f(jnp.ones(())) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "1.0 cannot be 1.0") @parameterized.named_parameters(("with_jit", True), ("without_jit", False)) def test_debug_check_disabled_errors(self, with_jit): def f(x): checkify.debug_check(jnp.all(x != x), "{x} cannot be {x}", x=x) return x checked_f = checkify.checkify(f, errors={}) if with_jit: checked_f = jax.jit(checked_f) err, _ = checked_f(jnp.ones((1,))) self.assertIsNone(err.get()) def test_debug_check_jaxpr_roundtrip(self): def f(x): checkify.debug_check(jnp.all(x != x), "{x} cannot be {x}", x=x) return x x = jnp.ones(()) jaxpr = jax.make_jaxpr(f)(x) roundtrip_f = partial(core.eval_jaxpr, jaxpr, jaxpr.consts) checked_f = checkify.checkify(jax.jit(roundtrip_f)) err, _ = checked_f(jnp.ones(())) self.assertIsNotNone(err.get()) self.assertStartsWith(err.get(), "1.0 cannot be 1.0") def test_fmt_args_array_type_error(self): args_error = lambda: checkify.check(False, "{} world", "hello") with self.assertRaisesRegex(TypeError, "Formatting arguments"): checkify.checkify(args_error)() kwargs_error = lambda: checkify.check(False, "{hello} world", hello="hello") with self.assertRaisesRegex(TypeError, "Formatting arguments"): checkify.checkify(kwargs_error)() np_arrays_ok = lambda: checkify.check(False, "{} world", np.array(1.)) checkify.checkify(np_arrays_ok)() trees_ok = lambda: checkify.check(False, "{}", {"hello": jnp.array(1.)}) checkify.checkify(trees_ok)() def test_checkify_non_jax_type_input(self): _ = checkify.checkify(lambda x: 1.)("hi") # does not crash def test_checkify_static_args(self): @checkify.checkify def f(x): if x: return _ = jax.jit(f, static_argnums=(0,))(True) def test_check_pp_rule(self): jaxpr = jax.make_jaxpr(lambda: checkify.check(False, "hi"))() jaxpr.pretty_print(source_info=True, name_stack=True) # Does not crash. def test_checkify_trace_determinism(self): x = jnp.array(5.0) def make_fn(): def f(y): checkify.check(y > 0.0, "Value must be positive!") return y * 2.0 return checkify.checkify(f) lowered_1 = jax.jit(make_fn()).lower(x) lowered_2 = jax.jit(make_fn()).lower(x) self.assertEqual(lowered_1.as_text(), lowered_2.as_text()) def test_checkify_multithreaded_determinism(self): x = jnp.array(5.0) def trace_and_lower(): def f(y): checkify.check(y > 0.0, "Value must be positive!") checkify.check(y < 100.0, "Value must be less than 100!") return y * 2.0 return jax.jit(checkify.checkify(f)).lower(x).as_text() with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: futures = [executor.submit(trace_and_lower) for _ in range(16)] results = [f.result() for f in futures] for res in results[1:]: self.assertEqual(results[0], res) def test_checkify_inner_jit(self): @checkify.checkify def f(x): @jax.jit def inner(y): checkify.check(y > 0.0, "Inner check failed!") return y + 1.0 return inner(x) err, out = f(jnp.array(3.0)) self.assertEqual(out, 4.0) err, _ = f(jnp.array(-1.0)) with self.assertRaisesRegex(JaxRuntimeError, "Inner check failed!"): err.throw() def test_nested_checkify_namespace_collision(self): def g(y): checkify.check(y > 0, "inner check failed") return y def f(x, y): checkify.check(x > 0, "outer check failed") err, out = checkify.checkify(g)(y) checkify.check_error(err) return out checked_f = checkify.checkify(f) err, _ = checked_f(-1, 1) self.assertStartsWith(err.get(), "outer check failed") err, _ = checked_f(1, -1) self.assertStartsWith(err.get(), "inner check failed") def test_reinjected_error_with_disabled_checks(self): # A check under a disabled error category must not corrupt the code # namespace of the emitted error: re-injecting that error under an outer # checkify alongside further outer checks must keep messages intact. def inner(x): checkify.check(x < 100.0, "dropped check") # user_checks not enabled return jnp.log(x) inner_checked = checkify.checkify(inner, errors=checkify.float_checks) err_inner, _ = inner_checked(jnp.array(-1.0)) # NaN fires self.assertIsNotNone(err_inner.get()) self.assertStartsWith(err_inner.get(), "nan generated by primitive: log") def outer(err, x): checkify.check_error(err) checkify.check(x < 1000.0, "outer check") # passes return x outer_checked = checkify.checkify( outer, errors=checkify.float_checks | checkify.user_checks) err_out, _ = outer_checked(err_inner, jnp.array(1.0)) self.assertIsNotNone(err_out.get()) self.assertStartsWith(err_out.get(), "nan generated by primitive: log") def test_reinjected_error_after_while_loop_cond_check(self): # The while-loop rule checkifies its cond jaxpr an extra time with the # error outputs dropped; the codes involved must not corrupt the namespace # of the emitted error under later re-injection. def inner(x): def cond(c): checkify.check(c < 100.0, "cond check") # passes return c < 3.0 out = jax.lax.while_loop(cond, lambda c: c + 1.0, x) checkify.check(x < 10.0, "post-while check") # fires return out err_inner, _ = checkify.checkify(inner)(jnp.array(20.0)) self.assertStartsWith(err_inner.get(), "post-while check") def outer(err, x): checkify.check_error(err) checkify.check(x < 1000.0, "outer check") # passes return x err_out, _ = checkify.checkify(outer)(err_inner, jnp.array(1.0)) self.assertIsNotNone(err_out.get()) self.assertStartsWith(err_out.get(), "post-while check") if __name__ == "__main__": absltest.main(testLoader=jtu.JaxTestLoader())