# 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. import collections import functools import logging import textwrap import threading import time import unittest from absl.testing import absltest, parameterized import jax from jax import lax from jax._src import ad_checkpoint from jax._src import debugging from jax._src import dispatch from jax._src import test_util as jtu from jax.sharding import PartitionSpec as P import jax.numpy as jnp import numpy as np try: import rich except ModuleNotFoundError: rich = None jax.config.parse_flags_with_absl() jtu.request_cpu_devices(2) debug_print = debugging.debug_print def _format_multiline(text): return textwrap.dedent(text).lstrip() class DummyDevice: def __init__(self, platform, id, device_kind): self.platform = platform self.id = id self.device_kind = device_kind class DebugCallbackTest(jtu.JaxTestCase): def tearDown(self): super().tearDown() dispatch.runtime_tokens.clear() def test_error_with_non_callable(self): with self.assertRaisesRegex(TypeError, "callable"): jax.debug.callback("this is not debug.print!") def test_ordered_callback_sequential_and_in_order(self): # Tests that ordered=True callbacks run sequentially and in program order. lock = threading.Lock() active = 0 max_active = 0 order = [] def cb(i, x): nonlocal active, max_active del x with lock: active += 1 max_active = max(max_active, active) order.append(i) time.sleep(0.002) with lock: active -= 1 n = 32 @jax.jit def f(x): for i in range(n): jax.debug.callback(functools.partial(cb, i), x, ordered=True) return x f(jnp.zeros(1024, jnp.float32)).block_until_ready() jax.effects_barrier() self.assertEqual(max_active, 1) self.assertEqual(order, list(range(n))) @jtu.skip_on_flag("jax_skip_slow_tests", True) @jtu.run_on_devices("cpu") def test_async_deadlock(self): # See https://github.com/jax-ml/jax/issues/25861 def print_it(i, maxiter): self.assertIsInstance(i, jax.Array) self.assertIsInstance(maxiter, jax.Array) return i == maxiter # Using JAX here causes deadlock with async dispatch def run(pos): maxiter = 1000 def cond(v): return v[0] < maxiter def step(v): i, pos = v jax.debug.callback(print_it, i + 1, maxiter) return i + 1, pos + 1 val = jnp.array(0), pos val = jax.lax.while_loop(cond, step, val) return val[1] n_samples = 30 inputs = 10 * jax.random.normal( jax.random.key(42), shape=(n_samples, 128, 128) ) def mean(forest): norm = 1.0 / len(forest) return norm * sum(forest) post_mean = mean(tuple(run(x) for x in inputs)) jax.block_until_ready(post_mean) # This shouldn't deadlock. @jtu.thread_unsafe_test_class() # printing isn't thread-safe class DebugPrintTest(jtu.JaxTestCase): def tearDown(self): super().tearDown() dispatch.runtime_tokens.clear() def test_simple_debug_print_works_in_eager_mode(self): def f(x): debug_print('x: {}', x) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") def test_static_args(self): @jax.jit def f(arr): jax.debug.print("arr {array}, dtype: {dtype}, arr {array2}", array=arr, dtype=arr.dtype, array2=arr) arr = jnp.array([1, 2, 3], dtype=jnp.float32) with jtu.capture_stdout() as output: f(arr) jax.effects_barrier() self.assertEqual( output(), "arr [1. 2. 3.], dtype: float32, arr [1. 2. 3.]\n") def test_debug_print_works_with_named_format_strings(self): def f(x): debug_print('x: {x}', x=x) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") def test_multiple_debug_prints_should_print_multiple_values(self): def f(x): debug_print('x: {x}', x=x) debug_print('y: {y}', y=x + 1) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2\ny: 3\n") def test_can_stage_out_debug_print(self): @jax.jit def f(x): debug_print('x: {x}', x=x) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") def test_can_stage_out_debug_print_with_formatting(self): @jax.jit def f(x): debug_print('x: {x:.2f}', x=x) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2.00\n") @jtu.device_supports_buffer_donation() def test_can_stage_out_debug_print_with_donate_argnums(self): def f(x, y): debug_print('x: {x}', x=x) return x + y f = jax.jit(f, donate_argnums=0) with jtu.capture_stdout() as output: f(2, 3) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") def test_can_stage_out_ordered_print(self): @jax.jit def f(x): debug_print('x: {x}', x=x, ordered=True) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") @jtu.device_supports_buffer_donation() def test_can_stage_out_ordered_print_with_donate_argnums(self): def f(x, y): debug_print('x: {x}', x=x, ordered=True) return x + y f = jax.jit(f, donate_argnums=0) with jtu.capture_stdout() as output: f(2, 3) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") @jtu.device_supports_buffer_donation() def test_can_stage_out_prints_with_donate_argnums(self): def f(x, y): debug_print('x: {x}', x=x, ordered=True) debug_print('x: {x}', x=x) return x + y f = jax.jit(f, donate_argnums=0) with jtu.capture_stdout() as output: f(2, 3) jax.effects_barrier() self.assertEqual(output(), "x: 2\nx: 2\n") def test_can_double_stage_out_ordered_print(self): @jax.jit @jax.jit def f(x): debug_print('x: {x}', x=x, ordered=True) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), "x: 2\n") def test_can_stage_out_ordered_print_with_pytree(self): @jax.jit def f(x): struct = dict(foo=x) debug_print('x: {}', struct, ordered=True) with jtu.capture_stdout() as output: f(np.array(2, np.int32)) jax.effects_barrier() self.assertEqual(output(), f"x: {str(dict(foo=jnp.array(2, np.int32)))}\n") def test_debug_print_should_use_default_layout(self): data = np.array( [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14]], dtype=np.int32) @jax.jit def f(x): jax.debug.print("{}", x) with jtu.capture_stdout() as output: f(data) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" [[ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14] [ 1 2 3 4 5 6 7 8 9 10 12 13 14]] """)) def test_debug_print_respects_numpy_printoptions(self): def f(x): with np.printoptions(precision=2, suppress=True): jax.debug.print("{}", x) x = np.array([1.2345, 2.3456, 1E-7]) # Default numpy print options: with jtu.capture_stdout() as output: jax.debug.print("{}", x) self.assertEqual(output(), "[1.2345e+00 2.3456e+00 1.0000e-07]\n") # Modified print options without JIT: with jtu.capture_stdout() as output: f(x) jax.effects_barrier() self.assertEqual(output(), "[1.23 2.35 0. ]\n") # Modified print options with JIT: with jtu.capture_stdout() as output: jax.jit(f)(x) jax.effects_barrier() self.assertEqual(output(), "[1.23 2.35 0. ]\n") @parameterized.parameters([False, True]) def test_debug_print_in_unrolled_loop(self, use_jit): def body(i, _): jax.debug.print("{}", i) if use_jit: body = jax.jit(body) @jax.jit def f(): return jax.lax.fori_loop(0, 4, body, None, unroll=2) with jtu.capture_stdout() as output: f() jax.effects_barrier() actual = tuple(sorted(map(int, output().splitlines()))) self.assertEqual(actual, tuple(range(4))) def test_debug_print_extended_dtype(self): def f(k): jax.debug.print("{}", k) with jtu.capture_stdout(): f(jax.random.key(0)) # doesn't crash jax.effects_barrier() @jtu.thread_unsafe_test_class() # printing isn't thread-safe class DebugPrintTransformationTest(jtu.JaxTestCase): def test_debug_print_batching(self): @jax.vmap def f(x): debug_print('hello: {}', x) with jtu.capture_stdout() as output: f(jnp.arange(2)) jax.effects_barrier() self.assertEqual(output(), "hello: 0\nhello: 1\n") def test_debug_print_batching_with_diff_axes(self): @functools.partial(jax.vmap, in_axes=(0, 1)) def f(x, y): debug_print('hello: {} {}', x, y) with jtu.capture_stdout() as output: f(jnp.arange(2), jnp.arange(2)[None]) jax.effects_barrier() self.assertEqual(output(), "hello: 0 [0]\nhello: 1 [1]\n") def tested_debug_print_with_nested_vmap(self): def f(x): debug_print('hello: {}', x) # Call with # [[0, 1], # [2, 3], # [4, 5]] with jtu.capture_stdout() as output: # Should print over 0-axis then 1-axis jax.vmap(jax.vmap(f))(jnp.arange(6).reshape((3, 2))) jax.effects_barrier() self.assertEqual( output(), "hello: 0\nhello: 2\nhello: 4\nhello: 1\nhello: 3\nhello: 5\n") with jtu.capture_stdout() as output: # Should print over 1-axis then 0-axis jax.vmap(jax.vmap(f, in_axes=0), in_axes=1)(jnp.arange(6).reshape((3, 2))) jax.effects_barrier() self.assertEqual( output(), "hello: 0\nhello: 1\nhello: 2\nhello: 3\nhello: 4\nhello: 5\n") def test_debug_print_jvp_rule(self): def f(x): debug_print('x: {}', x) with jtu.capture_stdout() as output: jax.jvp(f, (1.,), (1.,)) jax.effects_barrier() self.assertEqual(output(), "x: 1.0\n") def test_debug_print_vjp_rule(self): def f(x): debug_print('x: {}', x) with jtu.capture_stdout() as output: jax.vjp(f, 1.) jax.effects_barrier() self.assertEqual(output(), "x: 1.0\n") def test_debug_print_in_custom_jvp(self): @jax.custom_jvp def print_tangent(x): return x @print_tangent.defjvp def _(primals, tangents): (x,), (t,) = primals, tangents debug_print("x_tangent: {}", t) return x, t def f(x): x = jnp.square(x) x = print_tangent(x) return x with jtu.capture_stdout() as output: x = jnp.array(1., jnp.float32) jax.jvp(f, (x,), (x,)) jax.effects_barrier() expected = jnp.array(2., jnp.float32) self.assertEqual(output(), f"x_tangent: {expected}\n") @unittest.skip("doesn't work yet!") # TODO(mattjj,sharadmv) def test_debug_print_in_custom_jvp_linearize(self): @jax.custom_jvp def print_tangent(x): return x @print_tangent.defjvp def _(primals, tangents): (x,), (t,) = primals, tangents debug_print("x_tangent: {}", t) return x, t def f(x): x = jnp.sin(x) x = print_tangent(x) return x with jtu.capture_stdout() as output: x = jnp.array(1., jnp.float32) y, f_lin = jax.linearize(f, x) jax.effects_barrier() self.assertEqual(output(), "") with jtu.capture_stdout() as output: _ = f_lin(x) jax.effects_barrier() expected = jnp.cos(jnp.array(1., jnp.float32)) self.assertEqual(output(), f"x_tangent: {expected}\n") def test_debug_print_grad_with_custom_vjp_rule(self): @jax.custom_vjp def print_grad(x): return x def print_grad_fwd(x): return x, None def print_grad_bwd(_, x_grad): debug_print("x_grad: {}", x_grad) return (x_grad,) print_grad.defvjp(print_grad_fwd, print_grad_bwd) def f(x): debug_print("x: {}", x) x = print_grad(x) return jnp.square(x) with jtu.capture_stdout() as output: jax.grad(f)(jnp.array(1., jnp.float32)) jax.effects_barrier() expected = jnp.array(2., jnp.float32) self.assertEqual(output(), f"x: 1.0\nx_grad: {expected}\n") # mattjj was here # def test_debug_print_transpose_rule(self): # def f(x): # debug_print('should never be called: {}', x) # return x # with jtu.capture_stdout() as output: # jax.linear_transpose(f, 1.)(1.) # jax.effects_barrier() # self.assertEqual(output(), "") @jtu.sample_product(ordered=[False, True]) def test_remat_of_debug_print(self, ordered): def f_(x): y = ad_checkpoint.checkpoint_name(x + 1., "y") z = ad_checkpoint.checkpoint_name(y * 2., "z") debug_print('y: {}, z: {}', y, z, ordered=ordered) return ad_checkpoint.checkpoint_name(jnp.exp(z), "w") # Policy that saves nothing so everything gets rematerialized, including the # debug callback f = jax.checkpoint(f_, policy=ad_checkpoint.nothing_saveable) with jtu.capture_stdout() as output: jax.grad(f)(2.) jax.effects_barrier() # We expect the print to happen twice since it is rematerialized. self.assertEqual(output(), "y: 3.0, z: 6.0\n" * 2) # Policy that does not save `z` so we will need to rematerialize the print f = jax.checkpoint( f_, policy=ad_checkpoint.save_any_names_but_these("z")) with jtu.capture_stdout() as output: jax.grad(f)(2.) jax.effects_barrier() # We expect the print to happen twice since it is rematerialized. self.assertEqual(output(), "y: 3.0, z: 6.0\n" * 2) def test_debug_print_in_staged_out_custom_jvp(self): @jax.jit def f(x): @jax.custom_jvp def g(x): debug_print("hello: {x}", x=x) return x def g_jvp(primals, tangents): (x,), (t,) = primals, tangents debug_print("goodbye: {x} {t}", x=x, t=t) return x, t g.defjvp(g_jvp) return g(x) with jtu.capture_stdout() as output: f(2.) jax.effects_barrier() self.assertEqual(output(), "hello: 2.0\n") with jtu.capture_stdout() as output: jax.jvp(f, (2.,), (3.,)) jax.effects_barrier() self.assertEqual(output(), "goodbye: 2.0 3.0\n") def test_debug_print_in_staged_out_custom_vjp(self): @jax.jit def f(x): @jax.custom_vjp def g(x): debug_print("hello: {x}", x=x) return x def g_fwd(x): debug_print("hello fwd: {x}", x=x) return x, x def g_bwd(x, g): debug_print("hello bwd: {x} {g}", x=x, g=g) return (g,) g.defvjp(fwd=g_fwd, bwd=g_bwd) return g(x) with jtu.capture_stdout() as output: f(2.) jax.effects_barrier() self.assertEqual(output(), "hello: 2.0\n") with jtu.capture_stdout() as output: _, f_vjp = jax.vjp(f, 2.) jax.effects_barrier() self.assertEqual(output(), "hello fwd: 2.0\n") with jtu.capture_stdout() as output: f_vjp(3.0) jax.effects_barrier() self.assertEqual(output(), "hello bwd: 2.0 3.0\n") @jtu.thread_unsafe_test_class() # printing isn't thread-safe class DebugPrintControlFlowTest(jtu.JaxTestCase): def _assertLinesEqual(self, text1, text2): def _count(lines): return collections.Counter(lines) self.assertDictEqual(_count(text1.split("\n")), _count(text2.split("\n"))) @jtu.sample_product(ordered=[False, True]) def test_can_print_inside_scan(self, ordered): def f(xs): def _body(carry, x): debug_print("carry: {carry}, x: {x}", carry=carry, x=x, ordered=ordered) return carry + 1, x + 1 return lax.scan(_body, 2, xs) with jtu.capture_stdout() as output: f(jnp.arange(2)) jax.effects_barrier() self.assertEqual( output(), _format_multiline(""" carry: 2, x: 0 carry: 3, x: 1 """)) @jtu.sample_product(ordered=[False, True]) def test_can_print_inside_for_loop(self, ordered): def f(x): def _body(i, x): debug_print("i: {i}", i=i, ordered=ordered) debug_print("x: {x}", x=x, ordered=ordered) return x + 1 return lax.fori_loop(0, 5, _body, x) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() expected = _format_multiline(""" i: 0 x: 2 i: 1 x: 3 i: 2 x: 4 i: 3 x: 5 i: 4 x: 6 """) if ordered: self.assertEqual(output(), expected) else: self._assertLinesEqual(output(), expected) @jtu.sample_product(ordered=[False, True]) def test_can_print_inside_while_loop_body(self, ordered): def f(x): def _cond(x): return x < 10 def _body(x): debug_print("x: {x}", x=x, ordered=ordered) return x + 1 return lax.while_loop(_cond, _body, x) with jtu.capture_stdout() as output: f(5) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" x: 5 x: 6 x: 7 x: 8 x: 9 """)) @jtu.sample_product(ordered=[False, True]) def test_can_print_inside_while_loop_cond(self, ordered): def f(x): def _cond(x): debug_print("x: {x}", x=x, ordered=ordered) return x < 10 def _body(x): return x + 1 return lax.while_loop(_cond, _body, x) with jtu.capture_stdout() as output: f(5) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" x: 5 x: 6 x: 7 x: 8 x: 9 x: 10 """)) with jtu.capture_stdout() as output: f(10) jax.effects_barrier() # Should run the cond once self.assertEqual(output(), _format_multiline(""" x: 10 """)) @jtu.sample_product(ordered=[False, True]) def test_can_print_in_batched_while_cond(self, ordered): def f(x): def _cond(x): debug_print("x: {x}", x=x, ordered=ordered) return x < 5 def _body(x): return x + 1 return lax.while_loop(_cond, _body, x) with jtu.capture_stdout() as output: jax.vmap(f)(jnp.arange(2)) jax.effects_barrier() if ordered: expected = _format_multiline(""" x: 0 x: 1 x: 1 x: 2 x: 2 x: 3 x: 3 x: 4 x: 4 x: 5 x: 5 x: 6 """) self.assertEqual(output(), expected) else: # When the print is unordered, the `cond` is called an additional time # after the `_body` runs, so we get more prints. expected = _format_multiline(""" x: 0 x: 1 x: 0 x: 1 x: 1 x: 2 x: 1 x: 2 x: 2 x: 3 x: 2 x: 3 x: 3 x: 4 x: 3 x: 4 x: 4 x: 5 x: 4 x: 5 x: 5 x: 5 """) self._assertLinesEqual(output(), expected) @jtu.sample_product(ordered=[False, True]) def test_can_print_inside_cond(self, ordered): def f(x): def true_fun(x): debug_print("true: {}", x, ordered=ordered) return x def false_fun(x): debug_print("false: {}", x, ordered=ordered) return x return lax.cond(x < 5, true_fun, false_fun, x) with jtu.capture_stdout() as output: f(5) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" false: 5 """)) with jtu.capture_stdout() as output: f(4) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" true: 4 """)) @jtu.sample_product(ordered=[False, True]) def test_can_print_inside_switch(self, ordered): def f(x): def b1(x): debug_print("b1: {}", x, ordered=ordered) return x def b2(x): debug_print("b2: {}", x, ordered=ordered) return x def b3(x): debug_print("b3: {}", x, ordered=ordered) return x return lax.switch(x, (b1, b2, b3), x) with jtu.capture_stdout() as output: f(0) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" b1: 0 """)) with jtu.capture_stdout() as output: f(1) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" b2: 1 """)) with jtu.capture_stdout() as output: f(2) jax.effects_barrier() self.assertEqual(output(), _format_multiline(""" b3: 2 """)) @jtu.thread_unsafe_test_class() # printing isn't thread-safe class DebugPrintParallelTest(jtu.JaxTestCase): def _assertLinesEqual(self, text1, text2): def _count(lines): return collections.Counter(lines) self.assertDictEqual(_count(text1.split("\n")), _count(text2.split("\n"))) def test_ordered_print_not_supported_in_pmap(self): @jax.pmap def f(x): debug_print("{}", x, ordered=True) if jax.device_count() == 1: self.skipTest("This test won't raise with 1 device.") if jtu.device_under_test() == "gpu": self.skipTest("Test does not raise under GPU.") if jtu.device_under_test() == "tpu" and jtu.get_tpu_version() > 3: self.skipTest("Test does not raise under TPU v4+.") regex = ( "The following ordered effects are not supported for more than 1" " device:*" ) with self.assertRaisesRegex(ValueError, regex): f(jnp.arange(jax.local_device_count())) def test_unordered_print_works_in_pmap(self): if jax.device_count() < 2: raise unittest.SkipTest("Test requires >= 2 devices.") @jax.pmap def f(x): debug_print("hello: {}", x, ordered=False) with jtu.capture_stdout() as output: f(jnp.arange(jax.local_device_count())) jax.effects_barrier() lines = [f"hello: {i}\n" for i in range(jax.local_device_count())] self._assertLinesEqual(output(), "".join(lines)) @jax.pmap def f2(x): debug_print('hello: {}', x) debug_print('hello: {}', x + 2) with jtu.capture_stdout() as output: f2(jnp.arange(2)) jax.effects_barrier() self._assertLinesEqual(output(), "hello: 0\nhello: 1\nhello: 2\nhello: 3\n") def test_unordered_print_with_jit(self): def f(x): debug_print("{}", x, ordered=False) return x mesh = jax.sharding.Mesh(np.array(jax.devices()), ['dev']) spec = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('dev')) out_spec = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) f = jax.jit(f, in_shardings=spec, out_shardings=out_spec) with jax.set_mesh(mesh): with jtu.capture_stdout() as output: f(np.arange(8, dtype=jnp.int32)) jax.effects_barrier() self.assertEqual(output(), "[0 1 2 3 4 5 6 7]\n") def f2(x): y = x.dot(x) debug_print("{}", y, ordered=False) return y f2 = jax.jit(f2, in_shardings=spec, out_shardings=out_spec) with jax.set_mesh(mesh): with jtu.capture_stdout() as output: f2(np.arange(8, dtype=jnp.int32)) jax.effects_barrier() self.assertEqual(output(), "140\n") def test_nested_jit_debug_print(self): def f(x): debug_print("{}", x) return x with jtu.capture_stdout() as output: jax.jit(jax.jit(f))(jnp.arange(8)) jax.effects_barrier() self.assertEqual(output(), "[0 1 2 3 4 5 6 7]\n") def test_unordered_print_of_jit_of_while(self): def f(x): def cond(carry): i, *_ = carry return i < 5 def body(carry): i, x = carry debug_print("{}", x, ordered=False) x = x + 1 return (i + 1, x) return lax.while_loop(cond, body, (0, x))[1] mesh = jax.sharding.Mesh(np.array(jax.devices()), ['dev']) spec = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('dev')) f = jax.jit(f, in_shardings=spec, out_shardings=spec) with jax.set_mesh(mesh): with jtu.capture_stdout() as output: f(np.arange(8, dtype=jnp.int32)) jax.effects_barrier() self.assertEqual(output(), "[0 1 2 3 4 5 6 7]\n" "[1 2 3 4 5 6 7 8]\n" "[2 3 4 5 6 7 8 9]\n" "[ 3 4 5 6 7 8 9 10]\n" "[ 4 5 6 7 8 9 10 11]\n") def test_unordered_print_works_in_pmap_of_while(self): if jax.device_count() < 2: raise unittest.SkipTest("Test requires >= 2 devices.") @jax.pmap def f(x): def cond(x): return x < 3 def body(x): debug_print("hello: {}", x, ordered=False) return x + 1 return lax.while_loop(cond, body, x) with jtu.capture_stdout() as output: f(jnp.arange(2)) jax.effects_barrier() self._assertLinesEqual( output(), "hello: 0\nhello: 1\nhello: 2\n" "hello: 1\nhello: 2\n") def test_incorrectly_formatted_string(self): @jax.jit def f(x): debug_print("hello: {x}", x) return x with self.assertRaises(KeyError): f(jnp.arange(2)) jax.effects_barrier() @jax.jit def f(x): debug_print("hello: {}", x=x) return x with self.assertRaises(IndexError): f(jnp.arange(2)) jax.effects_barrier() def test_format_string_errors_with_unused_args(self): @jax.jit def f(x): debug_print("hello: {x}", x=x, y=x) return x with self.assertRaisesRegex(ValueError, "Unused keyword arguments"): f(jnp.arange(2)) jax.effects_barrier() @jax.jit def g(x): debug_print("hello", x) return x with self.assertRaisesRegex(ValueError, "Unused positional arguments"): g(jnp.arange(2)) jax.effects_barrier() def test_accidental_fstring(self): @jax.jit def f(x): debug_print(f"hello: {x}", x=x) return x with self.assertRaisesRegex(ValueError, "You may be passing an f-string"): f(jnp.arange(2)) jax.effects_barrier() @jtu.thread_unsafe_test_class() # logging isn't thread-safe class VisualizeShardingTest(jtu.JaxTestCase): def _create_devices(self, shape): num_devices = np.prod(shape) devices = [DummyDevice("CPU", i, "CPU") for i in range(num_devices)] return np.array(devices).reshape(shape) @parameterized.parameters( (1, 1), (1, 2), (1, 5), (2, 2), (2, 3), (2, 4), (2, 5)) def test_make_color_iter_uses_all_colors(self, num_rows, num_cols): num_colors = num_rows * num_cols colors = list(debugging.make_color_iter(lambda x: x, num_rows, num_cols)) self.assertLen(colors, num_colors) self.assertArraysEqual(np.sort(colors), np.linspace(0, 1, num_colors)) def test_trivial_sharding(self): mesh = jax.sharding.Mesh(self._create_devices(1), ['x']) pspec = jax.sharding.PartitionSpec('x') sd = jax.sharding.NamedSharding(mesh, pspec) shape = (5,) with jtu.capture_stdout() as output: debugging.visualize_sharding(shape, sd) self.assertEqual(output(), _format_multiline(""" ┌───────┐ │ CPU 0 │ └───────┘ """)) def test_trivial_sharding_with_scale(self): mesh = jax.sharding.Mesh(self._create_devices(1), ['x']) pspec = jax.sharding.PartitionSpec('x') sd = jax.sharding.NamedSharding(mesh, pspec) shape = (5,) with jtu.capture_stdout() as output: debugging.visualize_sharding(shape, sd, scale=8.) self.assertEqual(output(), _format_multiline(""" ┌──────────────────────────────────────┐ │ CPU 0 │ └──────────────────────────────────────┘ """)) def test_full_sharding(self): mesh = jax.sharding.Mesh(self._create_devices((8, 4)), ['x', 'y']) pspec = jax.sharding.PartitionSpec('x', 'y') sd = jax.sharding.NamedSharding(mesh, pspec) shape = (8, 8) with jtu.capture_stdout() as output: debugging.visualize_sharding(shape, sd) expected = _format_multiline(""" ┌───────┬───────┬───────┬───────┐ │ CPU 0 │ CPU 1 │ CPU 2 │ CPU 3 │ ├───────┼───────┼───────┼───────┤ │ CPU 4 │ CPU 5 │ CPU 6 │ CPU 7 │ ├───────┼───────┼───────┼───────┤ │ CPU 8 │ CPU 9 │CPU 10 │CPU 11 │ ├───────┼───────┼───────┼───────┤ │CPU 12 │CPU 13 │CPU 14 │CPU 15 │ ├───────┼───────┼───────┼───────┤ │CPU 16 │CPU 17 │CPU 18 │CPU 19 │ ├───────┼───────┼───────┼───────┤ │CPU 20 │CPU 21 │CPU 22 │CPU 23 │ ├───────┼───────┼───────┼───────┤ │CPU 24 │CPU 25 │CPU 26 │CPU 27 │ ├───────┼───────┼───────┼───────┤ │CPU 28 │CPU 29 │CPU 30 │CPU 31 │ └───────┴───────┴───────┴───────┘ """) self.assertEqual(output(), expected) def test_sharding_with_replication(self): shape = (8, 8) mesh = jax.sharding.Mesh(self._create_devices((8, 4)), ['x', 'y']) pspec = jax.sharding.PartitionSpec('x', None) sd = jax.sharding.NamedSharding(mesh, pspec) with jtu.capture_stdout() as output: debugging.visualize_sharding(shape, sd) expected = _format_multiline(""" ┌───────────────────────┐ │ CPU 0,1,2,3 │ ├───────────────────────┤ │ CPU 4,5,6,7 │ ├───────────────────────┤ │ CPU 8,9,10,11 │ ├───────────────────────┤ │ CPU 12,13,14,15 │ ├───────────────────────┤ │ CPU 16,17,18,19 │ ├───────────────────────┤ │ CPU 20,21,22,23 │ ├───────────────────────┤ │ CPU 24,25,26,27 │ ├───────────────────────┤ │ CPU 28,29,30,31 │ └───────────────────────┘ """) self.assertEqual(output(), expected) mesh = jax.sharding.Mesh(self._create_devices((4, 2)), ['x', 'y']) pspec = jax.sharding.PartitionSpec(None, 'y') sd = jax.sharding.NamedSharding(mesh, pspec) with jtu.capture_stdout() as output: debugging.visualize_sharding(shape, sd) expected = _format_multiline(""" ┌───────────┬───────────┐ │ │ │ │ │ │ │ │ │ │ │ │ │CPU 0,2,4,6│CPU 1,3,5,7│ │ │ │ │ │ │ │ │ │ │ │ │ └───────────┴───────────┘ """) self.assertEqual(output(), expected) def test_visualize_wide_array(self): shape = (128, 10000) mesh = jax.sharding.Mesh(self._create_devices((8, 4)), ['x', 'y']) pspec = jax.sharding.PartitionSpec('x', None) sd = jax.sharding.NamedSharding(mesh, pspec) with jtu.capture_stdout() as output: debugging.visualize_sharding(shape, sd) expected = _format_multiline(""" ┌──────────────────────────────────────────────────────────────────────────────┐ │ CPU 0,1,2,3 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 4,5,6,7 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 8,9,10,11 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 12,13,14,15 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 16,17,18,19 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 20,21,22,23 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 24,25,26,27 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ CPU 28,29,30,31 │ └──────────────────────────────────────────────────────────────────────────────┘ """) self.assertEqual(output(), expected) def test_visualize_sharding_shard_map(self): mesh = jtu.create_mesh((2,), 'x') def f(): a = jnp.zeros(1000) debugging.visualize_array_sharding(a) return a with jtu.capture_stdout(): f() # doesn't crash with jtu.capture_stdout(): jax.jit(f, out_shardings=jax.NamedSharding(mesh, P('x')))() # doesn't crash with jtu.capture_stdout(): jax.shard_map(f, mesh=mesh, in_specs=P(None), out_specs=P("x"))() # doesn't crash with jtu.capture_stdout(): jax.shard_map(f, mesh=mesh, in_specs=P(None), out_specs=P("x"), check_vma=False)() # doesn't crash class InspectShardingTest(jtu.JaxTestCase): def test_inspect_sharding_is_called_in_jit_sharded(self): if jtu.is_cloud_tpu(): raise unittest.SkipTest("Inspect sharding is not supported on libtpu.") is_called = False def _cb(sd): nonlocal is_called is_called = True self.assertIsInstance(sd, jax.sharding.Sharding) self.assertLen(sd.device_set, len(jax.devices())) def f(x): debugging.inspect_array_sharding(x, callback=_cb) return jnp.square(x) mesh = jax.sharding.Mesh(np.array(jax.devices()), ['dev']) spec = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('dev')) out_spec = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) f = jax.jit(f, in_shardings=spec, out_shardings=out_spec) with jax.set_mesh(mesh): f(np.arange(8, dtype=jnp.int32)) self.assertTrue(is_called) def test_inspect_sharding_is_called_in_jit(self): is_called = False def _cb(sd): nonlocal is_called is_called = True self.assertIsInstance(sd, jax.sharding.Sharding) self.assertLen(sd.device_set, 1) def f_(x): debugging.inspect_array_sharding(x, callback=_cb) return jnp.square(x) f = jax.jit(f_) f(np.arange(8, dtype=jnp.int32)) self.assertTrue(is_called) # Test in grad is_called = False f = jax.jit(jax.grad(lambda x: f_(x).sum())) f(np.arange(8, dtype=jnp.float32)) self.assertTrue(is_called) def test_inspect_sharding_3d_jit(self): def _cb(sd): self.assertIsInstance(sd, jax.sharding.NamedSharding) self.assertLen(sd.device_set, 2) def f_(x): debugging.inspect_array_sharding(x, callback=_cb) return jnp.square(x) f = jax.jit(f_) mesh = jtu.create_mesh((2,), ('x')) s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('x')) arr = jax.device_put(np.arange(8).reshape(2, 2, 2), s) f(arr) def _get_output_set(output, num_lines): """Return a set of strings where each string is num_lines.""" output = output().strip().split("\n") return { "\n".join(output[i : i + num_lines]) for i in range(0, len(output), num_lines) } @jtu.thread_unsafe_test_class() # printing isn't thread-safe class PartitionedDebugCallbackTest(jtu.JaxTestCase): def setUp(self): super().setUp() if (jtu.device_under_test() not in ("cpu", "gpu")): raise unittest.SkipTest( f"Test requires CPU or GPU devices. Got {jtu.device_under_test()}" ) if len(jax.devices()) < 2: raise unittest.SkipTest("Test requires >= 2 devices.") def tearDown(self): super().tearDown() dispatch.runtime_tokens.clear() def test_partitioned_debug_callback(self): def f_(x): debug_print("hello: {x}", x=x, partitioned=True) f = jax.jit(f_) mesh = jtu.create_mesh((1, 1, 2,), ("x", "y", "z")) s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("x", "y", "z")) arr = jax.device_put(np.arange(24).reshape(2, 3, 4), s) with jtu.capture_stdout() as output: with jax.set_mesh(mesh): f(arr) jax.effects_barrier() expected = { _format_multiline(""" hello: [[[ 0 1] [ 4 5] [ 8 9]] [[12 13] [16 17] [20 21]]]"""), _format_multiline(""" hello: [[[ 2 3] [ 6 7] [10 11]] [[14 15] [18 19] [22 23]]]"""), } self.assertEqual(_get_output_set(output, 7), expected) def test_partitioned_debug_callback_compute(self): def f(x): debug_print("hello: {x}", x=x.sum(), partitioned=True) mesh = jtu.create_mesh((2,), ("x",)) arr = jax.device_put(np.arange(8), jax.NamedSharding(mesh, jax.P("x"))) with jtu.capture_stdout(): with jax.set_mesh(mesh): f(arr) jax.effects_barrier() def test_debug_print_batching(self): @jax.vmap def f_(x): debug_print("hello: {}", x, partitioned=True) f = jax.jit(f_) mesh = jtu.create_mesh((1, 1, 2), ("x", "y", "z")) s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("x", "y", "z")) arr = np.arange(24).reshape(2, 3, 4) arr = jax.device_put(arr, s) with jtu.capture_stdout() as output: with jax.set_mesh(mesh): f(arr) jax.effects_barrier() expected = { _format_multiline(""" hello: [[0 1] [4 5] [8 9]]"""), _format_multiline(""" hello: [[ 2 3] [ 6 7] [10 11]]"""), _format_multiline(""" hello: [[14 15] [18 19] [22 23]]"""), _format_multiline(""" hello: [[12 13] [16 17] [20 21]]"""), } self.assertEqual(_get_output_set(output, 3), expected) def test_debug_print_batching_with_diff_axes(self): @functools.partial(jax.vmap, in_axes=(0, 1)) def f_(x, y): debug_print("hello: {} {}", x, y, partitioned=True) f = jax.jit(f_) mesh = jtu.create_mesh((2,), ("x")) s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("x")) x = np.arange(4).reshape(2, 2) x = jax.device_put(x, s) y = np.arange(4).reshape(2, 2) + 6 y = jax.device_put(y, s) with jtu.capture_stdout() as output: with jax.set_mesh(mesh): f(x, y) jax.effects_barrier() expected = { "hello: [2 3] [9]", "hello: [0 1] [6]", "hello: [0 1] [8]", "hello: [2 3] [7]", } self.assertEqual(_get_output_set(output, 1), expected) def test_debug_print_with_logging(self): logger_name = "jax._src.debugging" jax_logger = logging.getLogger(logger_name) class RecordHandler(logging.Handler): def __init__(self): logging.Handler.__init__(self) self.records = [] def emit(self, record): self.records.append(record) record_handler = RecordHandler() jax_logger.handlers.append(record_handler) def log_fn(x): x = x * x jax.debug.log("x={}", x) return x * x self.assertEqual(jax.jit(log_fn)(2), 16) jax_logger.removeHandler(record_handler) self.assertEqual(record_handler.records[0].msg, "x=4") def test_debug_print_with_nested_vmap(self): @jax.vmap @jax.vmap def f_(x): debug_print("hello: {}", x, partitioned=True) f = jax.jit(f_) mesh = jtu.create_mesh((1, 1, 2), ("x", "y", "z")) s = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("x", "y", "z")) arr = np.arange(24).reshape(2, 3, 4) arr = jax.device_put(arr, s) with jtu.capture_stdout() as output: with jax.set_mesh(mesh): f(arr) jax.effects_barrier() expected = { "hello: [14 15]", "hello: [12 13]", "hello: [18 19]", "hello: [16 17]", "hello: [22 23]", "hello: [20 21]", "hello: [2 3]", "hello: [0 1]", "hello: [6 7]", "hello: [10 11]", "hello: [4 5]", "hello: [8 9]", } self.assertEqual(_get_output_set(output, 1), expected) class TwoCallPatternTest(jtu.JaxTestCase): def test_debug_print_two_call_form(self): def f(x): jax.debug.print(ordered=True)("x: {x}", x=x) with jtu.capture_stdout() as output: f(42) jax.effects_barrier() self.assertEqual(output(), "x: 42\n") def test_debug_callback_two_call_form(self): results = [] def cb(val): results.append(val) def f(x): jax.debug.callback(ordered=True)(cb, x) f(100) jax.effects_barrier() self.assertEqual(results, [100]) def test_option_as_regular_kwarg_in_second_call(self): # partitioned=True is treated purely as a formatting parameter, so it raises ValueError since it's unused. with self.assertRaisesRegex( ValueError, "Unused keyword arguments to `jax.debug.print`: \\['partitioned'\\]", ): jax.debug.print(ordered=True)("hello {a}", a=42, partitioned=True) def test_debug_callback_unexpected_arguments(self): with self.assertRaisesRegex( TypeError, "debug_callback received unexpected arguments in the two-call form", ): jax.debug.callback(ordered=True, a=42)(lambda x, a: x * a) with self.assertRaisesRegex( TypeError, "debug_callback received unexpected arguments in the two-call form", ): jax.debug.callback(None, 42, ordered=True)(lambda x, a: x * a) def test_debug_print_unexpected_arguments(self): with self.assertRaisesRegex( TypeError, "debug_print received unexpected arguments in the two-call form", ): jax.debug.print(ordered=True, a=42)("hello {a}") with self.assertRaisesRegex( TypeError, "debug_print received unexpected arguments in the two-call form", ): jax.debug.print(None, 42, ordered=True)("hello {a}") if not rich: del VisualizeShardingTest if __name__ == '__main__': absltest.main(testLoader=jtu.JaxTestLoader())