rocm_jax/tests/host_callback_test.py

# Copyright 2020 Google LLC
#
# 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 functools
import itertools
import logging
import os
import re
import threading
import time
from typing import Callable, Optional, Sequence
from unittest import SkipTest

from absl.testing import absltest
from absl.testing import parameterized

from jax import api
from jax import lax
from jax import numpy as jnp
from jax import test_util as jtu
from jax.config import config
from jax.experimental import host_callback as hcb
from jax.lib import xla_bridge
from jax.util import prod
import numpy as np

config.parse_flags_with_absl()
FLAGS = config.FLAGS


class _TestingOutputStream(object):
  """Use as `output_stream` for tests."""

  def __init__(self):
    self._output = []
    self.test_method_name = None

  def write(self, what: str) -> None:
    print(f"output_stream[{self.test_method_name}]: {what}", end="")
    self._output.append(what)

  @property
  def output(self):
    return "".join(self._output)

  @property
  def output_sorted_by_device(self):
    # Assume that the output is a sequence of strings including metadata
    # and data, with metadata containing `device: xxx`
    by_device = []  # each element is a pair (device, str_list)
    for s in self._output:
      m = re.match(r'.*device: (\S+)', s)
      if m:
        by_device.append((m.group(1), []))
      by_device[-1][1].append(s)

    sorted_by_device = sorted(by_device, key=lambda x: x[0])
    return "\n".join(itertools.chain(*[s[1] for s in sorted_by_device]))

  def __str__(self):
    return "TestingOutputStream"

  def reset(self):
    self._output = []


testing_stream = _TestingOutputStream()


def fun1(a):
  y = hcb.id_print(a * 2., what="a * 2", output_stream=testing_stream)
  y = hcb.id_print(y * 3., what="y * 3", output_stream=testing_stream, result=y)
  return y**2  # Some computation to make the gradient interesting


def fun1_equiv(a):  # Numerical equivalent of fun`
  return (a * 2.)**2

def maybe_print(do_print: bool, arg, what: str, tap_with_device: Optional[bool] = False):
  """Conditionally print on testing_string"""
  if do_print:
    return hcb.id_print(arg, what=what,
                        output_stream=testing_stream, tap_with_device=tap_with_device)
  else:
    return arg

ignore_jit_of_pmap_warning = functools.partial(
  jtu.ignore_warning, message=".*jit-of-pmap.*")

def assertMultiLineStrippedEqual(tst: jtu.JaxTestCase,
                                 expected: str, what: str):
  """A variant that preprocesses the string to eliminate non-determinism in
  floating point values, and several uninteresting id_tap primitive params.
  """
  # Sometimes we get floating points in the output; we round them
  def repl_floats(match_group):
    matched = match_group.group(0)
    if matched == ".": return matched
    x = np.around(float(matched), decimals=2)
    return f"{x:.2f}"
  what = re.sub(r"\-?\d*\.[\-\def]*", repl_floats, what)
  what = re.sub(r"output_stream=[^\]\n,]*,?", "", what)
  what = re.sub(r"threshold=[^\]\n,]*,?", "", what)
  what = re.sub(r"bwd=[^\]\n]*", "", what)
  what = re.sub(r"out_trees=[^\]\n]*", "", what)
  what = re.sub(r"fwd_jaxpr_thunk=[^\]\n]*", "", what)
  what = re.sub(r"jvp_jaxpr_thunk=[^\]\n]*", "", what)
  # Empty lines
  what = re.sub(r"^\s*\n", "", what, flags=re.MULTILINE)
  def repl_func(match_group):
    matched = match_group.group(0)
    if "function _print_consumer" in matched:
      return "tap_func_=_print"
    else:
      return "..."
  what = re.sub(r"tap_func_=([^\]\n,]*),?", repl_func, what)
  tst.assertMultiLineStrippedEqual(expected, what)


prev_xla_flags = None
def setUpModule():
  global prev_xla_flags
  # This will control the CPU devices. On TPU we always have 2 devices
  prev_xla_flags = jtu.set_host_platform_device_count(2)


# Reset to previous configuration in case other test modules will be run.
def tearDownModule():
  prev_xla_flags()


def assertMultiDeviceOutputEqual(tst: jtu.JaxTestCase,
                                 expected_2CPUs: str):
  """Check that the multi-device output is equal to the expected.

  The tests run with 2 CPU devices on CPU (due to the flag), also
  on TPU (due to how the TPU tests are set up), but only 1 device on
  GPU. We adjust the expected output here for 1 device.

  Args:
    expected_2CPUs: the expected output for 2 CPUs. If there is only
      one device, this is trimmed to the first device. If the current
      device_under_test is not a CPU, then we change the names
  """
  assert api.device_count() in (1, 2)
  expected = expected_2CPUs
  if api.device_count() == 1:
    start_device_1 = expected.find('device: cpu:1')
    if start_device_1 >= 0:
      expected = expected[0:start_device_1]

  def replace_device_name(m) -> str:
    return str(api.devices()[int(m.group(1))])
  expected = re.sub(r'cpu:(\d+)', replace_device_name, expected)
  what = testing_stream.output_sorted_by_device
  return assertMultiLineStrippedEqual(tst, expected, what)




class HostCallbackTest(jtu.JaxTestCase):

  def setUp(self):
    testing_stream.reset()
    testing_stream.test_method_name = self._testMethodName
    self.old_flags = os.getenv("XLA_FLAGS", "")

  def tearDown(self) -> None:
    if os.getenv("XLA_FLAGS") != self.old_flags:
      os.environ["XLA_FLAGS"] = self.old_flags
      xla_bridge.get_backend.cache_clear()
    hcb.barrier_wait("HostCallbackTest.tearDown")

  def helper_set_hlo_dump(self):
    flags_str = os.getenv("XLA_FLAGS", "")
    os.environ["XLA_FLAGS"] = f"{flags_str} --xla_dump_to=/tmp/xla_dump"
    # Clear any cached backends so new CPU backend will pick up the env var.
    xla_bridge.get_backend.cache_clear()

  def test_eval(self):
    # TODO: renable jaxpr golden tests when changing host_callback
    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(fun1)(5.)))

    self.assertAllClose((5. * 2.) ** 2, fun1(5.))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        what: a * 2
        10.00
        what: y * 3
        30.00""", testing_stream.output)
    testing_stream.reset()

  def test_with_tuple_results(self):
    def func2(x):
      x1, y1 = hcb.id_print((x * 2., x * 3.), output_stream=testing_stream)
      return x1 + y1

    self.assertEqual(3. * (2. + 3.), func2(3.))
    hcb.barrier_wait()

    assertMultiLineStrippedEqual(self, """
        [ 6.00
          9.00 ]""", testing_stream.output)
    testing_stream.reset()

  def test_with_dict_results(self):
    def func2(x):
      res = hcb.id_print(dict(a=x * 2., b=x * 3.), output_stream=testing_stream)
      return res["a"] + res["b"]

    self.assertEqual(3. * (2. + 3.), func2(3.))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        { a=6.00
          b=9.00 }""", testing_stream.output)
    testing_stream.reset()

  def test_with_result(self):
    def func2(x):
      x1 = hcb.id_print((x * 2., x * 3.), result=x * 4.,
                        output_stream=testing_stream)
      return x1

    self.assertEqual(3. * 4., func2(3.))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        [ 6.00
          9.00 ]""", testing_stream.output)
    testing_stream.reset()

  def test_print_with_device(self):
    def func2(x):
      x1 = hcb.id_print((x * 2., x * 3.), result=x * 4.,
                        output_stream=testing_stream,
                        tap_with_device=True)
      return x1

    self.assertEqual(3. * 4., func2(3.))
    hcb.barrier_wait()
    assertMultiDeviceOutputEqual(self, """
      device: cpu:0
      [ 6.00
        9.00 ]""")
    testing_stream.reset()

  def test_eval_tap_exception(self):
    # Simulate a tap error
    def tap_err(*args, **kwargs):
      raise NotImplementedError

    def func(x):
      x1 = hcb.id_print(x + 1, what="x1", output_stream=testing_stream)
      x2 = hcb.id_tap(tap_err, x1 + 1)
      x3 = hcb.id_print(x2 + 1, what="x3", output_stream=testing_stream)
      return x3

    with self.assertRaises(hcb.TapFunctionException):
      func(0)
      hcb.barrier_wait()

    # We should have received everything before the error
    assertMultiLineStrippedEqual(self, """
        what: x1
        1
        what: x3
        3""", testing_stream.output)
    testing_stream.reset()

  def test_jit_simple(self):
    jit_fun1 = api.jit(lambda x: 3. * hcb.id_print(
        2. * x, what="here", output_stream=testing_stream))
    self.assertAllClose(6. * 5., jit_fun1(5.))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        what: here
        10.00""", testing_stream.output)
    testing_stream.reset()

  def test_jit_no_invars(self):
    def func():  # jitted function does not take arguments
      return hcb.id_print(42, output_stream=testing_stream)

    self.assertAllClose(42, api.jit(func)())
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
    42""", testing_stream.output)
    testing_stream.reset()

  def test_jit_multiple_invars(self):
    def func(x1, x2):
      return hcb.id_print(x1 + x2, output_stream=testing_stream)

    self.assertAllClose(42, api.jit(func)(40, 2))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
    42""", testing_stream.output)
    testing_stream.reset()

  def test_jit_constant(self):
    def func(x):
      return hcb.id_print(42, result=x, output_stream=testing_stream)

    self.assertAllClose(5, api.jit(func)(5))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
    42""", testing_stream.output)
    testing_stream.reset()

  def test_jit_sequence1(self):
    def func(x):
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      return hcb.id_print(x1 + 1, where="2", output_stream=testing_stream)

    logging.info("%s: %s", self._testMethodName,
          api.make_jaxpr(func)(1))
    logging.info("%s: %s", self._testMethodName,
          api.xla_computation(func)(1).as_hlo_text())
    self.assertEqual(2, api.jit(func)(1))
    hcb.barrier_wait()

    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: 2
        2""", testing_stream.output)
    testing_stream.reset()

  def test_jit2(self):
    """A sequence of JIT."""
    def func(x):
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      x2 = hcb.id_print(x1 + 1, where="2", output_stream=testing_stream)
      return x2

    self.assertEqual(2, api.jit(func)(1))
    self.assertEqual(11, api.jit(func)(10))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: 2
        2
        where: 1
        10
        where: 2
        11""", testing_stream.output)
    testing_stream.reset()

  def test_jit_result_unused(self):
    """We can id_print even if we don't use the result."""
    if not config.omnistaging_enabled:
      raise SkipTest("Test requires omnistaging")
    def func(x):
      hcb.id_print(x, where="1", output_stream=testing_stream)
      hcb.id_print(x + 1, where="2", output_stream=testing_stream)
      return x + 1

    self.assertEqual(2, api.jit(func)(1))
    self.assertEqual(11, api.jit(func)(10))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: 2
        2
        where: 1
        10
        where: 2
        11""", testing_stream.output)
    testing_stream.reset()

  def test_jit_nested(self):
    def func(x):
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      def func_nested(x):
        x2 = hcb.id_print(x + 1, where="nested", output_stream=testing_stream)
        return x2
      x3 = api.jit(func_nested)(x1)
      return hcb.id_print(x3 + 1, where="3", output_stream=testing_stream)

    self.assertEqual(3, api.jit(func)(1))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: nested
        2
        where: 3
        3""", testing_stream.output)
    testing_stream.reset()

  def test_jit_devices(self):
    """Running on multiple devices."""
    devices = api.local_devices()
    logging.info(f"{self._testMethodName}: has devices {devices}")
    def func(x, device_id):
      x1 = hcb.id_print(x, dev=str(device_id), output_stream=testing_stream)
      x2 = hcb.id_print(x1 + 1, dev=str(device_id), output_stream=testing_stream)
      return x2

    for d in devices:
      self.assertEqual(112, api.jit(func, device=d, static_argnums=1)(111, d.id))
    hcb.barrier_wait()
    logging.info(f"{self._testMethodName}: found output {testing_stream.output}")
    self.assertEqual(len(devices), len(re.findall(r"111", testing_stream.output)))
    self.assertEqual(len(devices), len(re.findall(r"112", testing_stream.output)))
    testing_stream.reset()

  @parameterized.named_parameters(
      jtu.cases_from_list(
          dict(
              testcase_name=f"_with_jit_{with_jit}",
              with_jit=with_jit)
          for with_jit in [True, False]))
  def test_pytree(self, with_jit=False):
    def func(x, what=""):
      """Returns some pytrees depending on x"""
      if what == "pair_1_x":
        return (1, x)
      elif what == "pair_x_2x":
        return (x, 2 * x)
      elif what == "dict":
        return dict(a=2 * x, b=3 * x)
      else:
        assert False
    tap_count = 0

    def tap_func(a, _, *, what=""):
      nonlocal tap_count
      tap_count += 1
      self.assertEqual(func(5, what), a)

    transform = api.jit if with_jit else lambda f: f
    for what in ("pair_1_x", "pair_x_2x", "dict"):
      transformed = transform(
          lambda x: hcb.id_tap(
              functools.partial(tap_func, what=what),
              func(x, what),
              result=func(x * 2, what))
          )(5)
      self.assertEqual(func(10, what), transformed)
    hcb.barrier_wait()  # Wait for receivers to be done
    self.assertEqual(3, tap_count)

  @parameterized.named_parameters(
      jtu.cases_from_list(
          dict(
              testcase_name=f"_concurrent_{concurrent}",
              concurrent=concurrent)
          for concurrent in [True, False]))
  def test_multiple_tap(self, concurrent=False):
    """Call id_tap multiple times, concurrently or in sequence. """
    if concurrent and jtu.device_under_test() == "gpu":
      # TODO(necula): it seems that on GPU if multiple host threads run
      # a jit computation, the multiple computations are interleaved on the
      # GPU. This can result in the outfeed trains being interleaved, which
      # will trigger an error. The solution is to fix on GPU the receiving
      # logic so that we can outfeed the train as one tuple, and receive it
      # one piece as a time. Then the trains should be atomic.
      # See also b/160692602.
      raise SkipTest("concurrent id_tap not supported on GPU")
    received = set()
    count = 5
    def pause_tap(idx, _):
      received.add(int(idx))
      logging.info(f"Starting do_tap {idx}. Sleeping 1sec ...")
      time.sleep(0.3)
      logging.info(f"Finish do_tap {idx}")

    def do_tap(idx):
      api.jit(lambda idx: hcb.id_tap(pause_tap, idx))(idx)

    if concurrent:
      threads = [
          threading.Thread(
              name=f"enqueue_tap_{idx}", target=do_tap, args=(idx,))
          for idx in range(count)
      ]
      [t.start() for t in threads]
      [t.join() for t in threads]
    else:
      for idx in range(count):
        do_tap(idx)

    hcb.barrier_wait()
    self.assertEqual(received, set(range(count)))

  # TODO(necula): see comment for test_multiple_tap.
  @jtu.skip_on_devices("gpu")
  def test_multiple_barriers(self):
    """Call barrier_wait concurrently."""

    def pause_tap(*args, **kwargs):
      logging.info("pause_tap waiting")
      time.sleep(0.3)
      logging.info("pause_tap done")

    def long_run(x):
      return hcb.id_tap(pause_tap, x)

    api.jit(long_run)(5.)

    def try_barrier(idx):
      logging.info(f"Starting test barrier {idx}")
      hcb.barrier_wait()
      logging.info(f"Finished test barrier {idx}")

    threads = [
        threading.Thread(
            name=f"barrier_{idx}", target=try_barrier, args=(idx,))
        for idx in range(3)
    ]
    [t.start() for t in threads]
    [t.join() for t in threads]

  @parameterized.named_parameters(
      jtu.cases_from_list(
          dict(
              testcase_name=f"_with_jit_{with_jit}",
              with_jit=with_jit)
          for with_jit in [True, False]))
  def test_cond(self, with_jit=False):
    """A conditional"""
    def func(x):
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      x2 = hcb.id_print(x1 + 1, where="2", output_stream=testing_stream)

      x4 = lax.cond(x % 2 == 0,
                    lambda x: hcb.id_print(x, where="cond_t",
                                           output_stream=testing_stream),
                    lambda x: hcb.id_print(-1, where="cond_f", result=x,
                                           output_stream=testing_stream),
                    x2 + 1)
      x5 = hcb.id_print(x4 + 1, where="end", output_stream=testing_stream)
      return x5

    transform = api.jit if with_jit else lambda f: f
    self.assertEqual(4, transform(func)(1))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: 2
        2
        where: cond_f
        -1
        where: end
        4""", testing_stream.output)
    testing_stream.reset()

  @parameterized.named_parameters(
      jtu.cases_from_list(
          dict(testcase_name=f"_with_jit_{with_jit}",
               with_jit=with_jit)
          for with_jit in [True, False]))
  def test_while_cond(self, with_jit=False):
    def func(x):
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      x2 = hcb.id_print(x1 + 1, where="2", output_stream=testing_stream)
      def body(x):
        x3 = hcb.id_print(x, where="w_b_1", output_stream=testing_stream)
        x4 = lax.cond(x % 2 == 0,
                      lambda x: hcb.id_print(x, where="w_b_t",
                                             output_stream=testing_stream),
                      lambda x: hcb.id_print(-1, where="w_b_f",
                                             result=x, output_stream=testing_stream),
                      x3 + 1)
        return hcb.id_print(x4, where="w_b_2", output_stream=testing_stream)
      x10 = lax.while_loop(lambda x: x <= 3, body, x2)
      res = hcb.id_print(x10, where="end", output_stream=testing_stream)
      return res

    transform = api.jit if with_jit else lambda f: f
    self.assertEqual(4, transform(func)(1))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: 2
        2
        where: w_b_1
        2
        where: w_b_t
        3
        where: w_b_2
        3
        where: w_b_1
        3
        where: w_b_f
        -1
        where: w_b_2
        4
        where: end
        4""", testing_stream.output)
    testing_stream.reset()

  def test_jit_while_pred_tap(self):
    """While with printing in the conditional."""
    def func(x):
      x1 = hcb.id_print(x, where="1")
      x10 = lax.while_loop(lambda x: hcb.id_print(x < 3,
                                                  where="w_p",
                                                  output_stream=testing_stream),
                           lambda x: hcb.id_print(x + 1, where="w_b",
                                                  output_stream=testing_stream),
                           x1)
      res = hcb.id_print(x10, where="3", output_stream=testing_stream)
      return res

    self.assertEqual(3, api.jit(func)(1))
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self,
        """
        where: w_p
        True
        where: w_b
        2
        where: w_p
        True
        where: w_b
        3
        where: w_p
        False
        where: 3
        3""", testing_stream.output)
    testing_stream.reset()

  @parameterized.named_parameters(
      jtu.cases_from_list(
          dict(
              testcase_name=f"_with_jit_{with_jit}",
              with_jit=with_jit)
          for with_jit in [True, False]))
  def test_scan_cond(self, with_jit=False):
    def func(x):
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      x2 = hcb.id_print(x1 + 1, where="2", output_stream=testing_stream)

      def body(c, x):
        x3 = hcb.id_print(x, where="s_1", output_stream=testing_stream)
        x4 = lax.cond(x % 2 == 0,
                      lambda x: hcb.id_print(x, where="s_t", output_stream=testing_stream),
                      lambda x: hcb.id_print(-1, where="s_f", result=x, output_stream=testing_stream),
                      x3 + 1)
        return (c, hcb.id_print(x4, where="s_2", output_stream=testing_stream))

      _, x10 = lax.scan(body, x2, jnp.arange(3))
      res = hcb.id_print(x10, where="10", output_stream=testing_stream)
      return res

    if with_jit:
      func = api.jit(func)
    res = func(1)
    self.assertAllClose(jnp.array([1, 2, 3]), res)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        where: 1
        1
        where: 2
        2
        where: s_1
        0
        where: s_t
        1
        where: s_2
        1
        where: s_1
        1
        where: s_f
        -1
        where: s_2
        2
        where: s_1
        2
        where: s_t
        3
        where: s_2
        3
        where: 10
        [1 2 3]""", testing_stream.output)
    testing_stream.reset()

  @parameterized.named_parameters(
      jtu.cases_from_list(
          dict(
              testcase_name=f"_shape_{shape}_dtype_{dtype}_nr_args={nr_args}",
              shape=shape,
              dtype=dtype,
              nr_args=nr_args) for nr_args in [1, 2]
          for shape in [(), (2,), (2, 3), (2, 3, 4)]
          for dtype in jtu.dtypes.all))
  def test_jit_types(self, nr_args=2, dtype=jnp.int16, shape=(2,)):
    if dtype in (jnp.complex64, jnp.complex128, jnp.bool_):
      raise SkipTest(f"id_print jit not implemented for {dtype}.")
    if jtu.device_under_test() == "tpu":
      if dtype in (jnp.int16,):
        raise SkipTest(f"transfering {dtype} not supported on TPU")
    args = [jnp.arange(prod(shape), dtype=dtype).reshape(shape)]
    if nr_args > 1:
      args = args * nr_args
    jit_fun1 = api.jit(lambda xs: hcb.id_print(
        xs,
        a_new_test="************",
        testcase_name=f"shape_{shape}_dtype_{dtype}_nr_args={nr_args}"))

    res = jit_fun1(args)
    self.assertAllClose(args, res)

  def test_jit_large(self):
    arg = jnp.arange(10000, dtype=jnp.int32).reshape((10, 10, 5, -1))
    api.jit(hcb.id_print)(arg)

  def test_jit_several_together(self):
    arg = jnp.arange(50, dtype=jnp.int32).reshape((10, 5))
    api.jit(lambda x, y: hcb.id_print((x, y, x * 2.)))(arg, jnp.ones(100, dtype=jnp.int32))

  def test_jit_interleaving(self):
    # Several jit's without data dependencies; they may interfere
    count = 0  # Count tap invocations
    nr_arrays = 5
    def tap_func(arg, _):
      nonlocal count
      assert len(arg) == nr_arrays
      count += 1
    # This is the function that we'll run multiple times
    def func(x, count):
      for i in range(count):
        x = hcb.id_tap(tap_func, [x + i for i in range(nr_arrays)])[-1]
      return x

    x = jnp.array(1, dtype=np.int32)
    res = 0
    for _ in range(10):
      # No dependencies between the jit invocations
      res += api.jit(lambda x: func(x, 10))(x)
    hcb.barrier_wait()
    self.assertEqual(100, count)

  def test_jit_tap_exception(self):
    # Simulate a tap error
    def tap_err(*args, **kwargs):
      raise NotImplementedError
    def func(x):
      x1 = hcb.id_print(x + 1, what="x1", output_stream=testing_stream)
      x2 = hcb.id_tap(tap_err, x1 + 1)
      x3 = hcb.id_print(x2 + 1, what="x3", output_stream=testing_stream)
      return x3

    res = api.jit(func)(0)  # No error yet
    with self.assertRaises(hcb.TapFunctionException):
      hcb.barrier_wait()

    # Even though the receiver thread raised, the main thread should still
    # return 3.
    self.assertEqual(3, res)
    # We should have received all others
    assertMultiLineStrippedEqual(self, """
        what: x1
        1
        what: x3
        3""", testing_stream.output)
    testing_stream.reset()

  def test_while(self):
    """Executing while, even without JIT uses compiled code"""
    y = jnp.ones(5)  # captured const

    def func(x):
      return lax.while_loop(
        lambda c: c[1] < 5,
        lambda c: (y, hcb.id_print(c[1], output_stream=testing_stream) + 1),
        (x, 1))
    func(y)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        1
        2
        3
        4""", testing_stream.output)
    testing_stream.reset()

  def test_jvp(self):
    jvp_fun1 = lambda x, xt: api.jvp(fun1, (x,), (xt,))
    res_primals, res_tangents = jvp_fun1(jnp.float32(5.), jnp.float32(0.1))
    self.assertAllClose(100., res_primals, check_dtypes=False)
    self.assertAllClose(4., res_tangents, check_dtypes=False)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        what: a * 2
        10.00
        transforms: ['jvp'] what: a * 2
        0.20
        what: y * 3
        30.00
        transforms: ['jvp'] what: y * 3
        0.60""", testing_stream.output)
    testing_stream.reset()

  def test_grad_primal_unused(self):
    if not config.omnistaging_enabled:
      raise SkipTest("Test requires omnistaging")
    # The output of id_print is not needed for backwards pass
    def func(x):
      return 2. * hcb.id_print(x * 3., what="x * 3",
                               output_stream=testing_stream)

    grad_func = api.grad(func)
    jaxpr = str(api.make_jaxpr(grad_func)(5.))
    # making the Jaxpr does not print anything
    hcb.barrier_wait()

    assertMultiLineStrippedEqual(self, """
        { lambda  ; a.
          let b = mul a 3.00
              c = id_tap[ arg_treedef_=*
                          nr_tapped_args_=1
                          tap_func_=_print   what='x * 3') ] b
              _ = mul c 2.00
              d = mul 1.00 2.00
              e _ = id_tap[ arg_treedef_=*
                            nr_tapped_args_=1
                            tap_func_=_print   what='x * 3')
                            transforms=(('jvp',), ('transpose',)) ] d 0.00
              f = mul e 3.00
          in (f,) }""", jaxpr)
    assertMultiLineStrippedEqual(self, "", testing_stream.output)
    testing_stream.reset()

    res_grad = grad_func(jnp.float32(5.))
    hcb.barrier_wait()

    self.assertAllClose(6., res_grad, check_dtypes=False)
    assertMultiLineStrippedEqual(self, """
        what: x * 3
        15.00
        transforms: ['jvp', 'transpose'] what: x * 3
        2.00""", testing_stream.output)
    testing_stream.reset()

  def test_grad_simple(self):
    def func(x):
      y = hcb.id_print(x * 2., what="x * 2", output_stream=testing_stream)
      return x * hcb.id_print(y * 3., what="y * 3",
                              output_stream=testing_stream)
    grad_func = api.grad(func)

    res_grad = grad_func(jnp.float32(5.))
    self.assertAllClose(2. * 5. * 6., res_grad, check_dtypes=False)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        what: x * 2
        10.00
        what: y * 3
        30.00
        transforms: ['jvp', 'transpose'] what: y * 3
        5.00
        transforms: ['jvp', 'transpose'] what: x * 2
        15.00""", testing_stream.output)
    testing_stream.reset()

  def test_grad_double(self):
    if not config.omnistaging_enabled:
      raise SkipTest("Test requires omnistaging")
    def func(x):
      y = hcb.id_print(x * 2., what="x * 2", output_stream=testing_stream)
      return x * (y * 3.)

    grad_func = api.grad(api.grad(func))
    # making the Jaxpr does not print anything
    _ = api.make_jaxpr(grad_func)(5.)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, "", testing_stream.output)

    res_grad = grad_func(jnp.float32(5.))

    self.assertAllClose(12., res_grad, check_dtypes=False)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        what: x * 2
        10.00
        transforms: ['jvp', 'transpose'] what: x * 2
        15.00
        transforms: ['jvp', 'transpose', 'jvp', 'transpose'] what: x * 2
        2.00
        transforms: ['jvp', 'transpose'] what: x * 2
        3.00""", testing_stream.output)
    testing_stream.reset()

  def test_vmap(self):
    vmap_fun1 = api.vmap(fun1)
    vargs = jnp.array([jnp.float32(4.), jnp.float32(5.)])
    vmap_fun1(vargs)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        transforms: [('batch', {'batch_dims': (0,)})] what: a * 2
        [ 8.00 10.00]
        transforms: [('batch', {'batch_dims': (0, 0)})] what: y * 3
        [24.00 30.00]""", testing_stream.output)
    testing_stream.reset()

  def test_vmap_not_batched(self):
    x = 3.
    def func(y):
      # x is not mapped, y is mapped
      _, y = hcb.id_print((x, y), output_stream=testing_stream)
      return x + y

    vmap_func = api.vmap(func)
    vargs = jnp.array([jnp.float32(4.), jnp.float32(5.)])
    _ = vmap_func(vargs)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        transforms: [('batch', {'batch_dims': (None, 0)})]
        [ 3.00
          [4.00 5.00] ]""", testing_stream.output)
    testing_stream.reset()

  def test_double_vmap(self):
    # A 2D tensor with x[i, j] = i + j using 2 vmap
    def sum(x, y):
      return hcb.id_print(x + y, output_stream=testing_stream)
    def sum_rows(xv, y):
      return api.vmap(sum, in_axes=(0, None))(xv, y)
    def sum_all(xv, yv):
      return api.vmap(sum_rows, in_axes=(None, 0))(xv, yv)

    xv = jnp.arange(5, dtype=np.int32)
    yv = jnp.arange(3, dtype=np.int32)
    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(sum_all)(xv, yv)))
    _ = sum_all(xv, yv)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        transforms: [('batch', {'batch_dims': (0,)}), ('batch', {'batch_dims': (0,)})]
        [[0 1 2 3 4]
        [1 2 3 4 5]
        [2 3 4 5 6]]""", testing_stream.output)
    testing_stream.reset()

  def test_vmap_while(self):
    """Vmap of while."""

    def func(x):
      # like max(x, 2)
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      x2 = lax.while_loop(lambda x: x < 2,
                           lambda x: hcb.id_print(x + 1, where="w_b",
                                                  output_stream=testing_stream),
                           x1)
      res = hcb.id_print(x2, where="3", output_stream=testing_stream)
      return res

    inputs = np.arange(5, dtype=np.int32)
    self.assertAllClose(np.array([2, 2, 2, 3, 4]), api.jit(api.vmap(func))(inputs),
                        check_dtypes=False)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        transforms: [('batch', {'batch_dims': (0,)})] where: 1
        [0 1 2 3 4]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_b
        [1 2 3 4 5]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_b
        [2 3 3 4 5]
        transforms: [('batch', {'batch_dims': (0,)})] where: 3
        [2 2 2 3 4]""", testing_stream.output)
    testing_stream.reset()

  def test_vmap_while_tap_cond(self):
    """Vmap of while, with a tap in the conditional."""

    def func(x):
      # like max(x, 2)
      x1 = hcb.id_print(x, where="1", output_stream=testing_stream)
      x2 = lax.while_loop(lambda x: hcb.id_print(x < 2, where="w_c",
                                                 output_stream=testing_stream),
                          lambda x: hcb.id_print(x + 1, where="w_b",
                                                 output_stream=testing_stream),
                          x1)
      res = hcb.id_print(x2, where="3", output_stream=testing_stream)
      return res

    inputs = np.arange(5, dtype=np.int32)
    res = api.jit(api.vmap(func))(inputs)
    hcb.barrier_wait()
    self.assertAllClose(np.array([2, 2, 2, 3, 4]), res, check_dtypes=False)
    assertMultiLineStrippedEqual(self, """
        transforms: [('batch', {'batch_dims': (0,)})] where: 1
        [0 1 2 3 4]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_c
        [ True  True False False False]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_b
        [1 2 3 4 5]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_c
        [ True False False False False]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_b
        [2 3 3 4 5]
        transforms: [('batch', {'batch_dims': (0,)})] where: w_c
        [False False False False False]
        transforms: [('batch', {'batch_dims': (0,)})] where: 3
        [2 2 2 3 4]""", testing_stream.output)
    testing_stream.reset()

  def test_pmap(self):
    xv = jnp.arange(api.device_count(), dtype=jnp.int32)
    def fun1(x, do_print=False): # x: i32
      return maybe_print(do_print, x * 2, "x * 2", tap_with_device=True)

    pmap_fun1 = api.pmap(functools.partial(fun1, do_print=True))
    res = pmap_fun1(xv)
    hcb.barrier_wait()
    expected_res = api.pmap(functools.partial(fun1, do_print=False))(xv)
    self.assertAllClose(expected_res, res, check_dtypes=False)
    # Assertion text is for 2 devices (also works for 1 device)
    assertMultiDeviceOutputEqual(self, """
        device: cpu:0 what: x * 2
        0
        device: cpu:1 what: x * 2
        2""")
    testing_stream.reset()

  def test_pmap_vmap(self):
    # A matrix M[ij] = i * 10 + j
    nr_devices = api.device_count()
    shape = (nr_devices, 3)
    matrix = np.fromfunction(lambda i, j: 10. * i + j, shape,
                             dtype=np.int32)

    def fun1(x, do_print=False):  # x: i32
      return maybe_print(do_print, x * 2, "x * 2", tap_with_device=True)

    pmap_vmap_fun1 = api.pmap(api.vmap(functools.partial(fun1, do_print=True)))

    res = pmap_vmap_fun1(matrix)
    hcb.barrier_wait()
    expected_res = api.pmap(api.vmap(functools.partial(fun1, do_print=False)))(matrix)
    self.assertAllClose(expected_res, res, check_dtypes=False)
    # Assertion text is for 2 devices (also works for 1 device)
    assertMultiDeviceOutputEqual(self, """
        device: cpu:0 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
        [0.00 2.00 4.00]
        device: cpu:1 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
        [20.00 22.00 24.00]""")
    testing_stream.reset()

  def test_pmap_pmap_vmap(self):
    # A matrix M[ijk] = i * 100 + j * 10 + k
    nr_devices = api.local_device_count()
    if nr_devices % 2 != 0:
      raise SkipTest("test works only on even number of devices")

    shape = (2, nr_devices // 2, 3)
    matrix = np.fromfunction(lambda i, j, k: 100. * i + 10. * j + k, shape,
                             dtype=np.float32)
    def fun1(x, do_print=False): # x: f32
      y = maybe_print(do_print, x * 2., "x * 2", tap_with_device=True)
      return y ** 2

    pmap_fun1 = api.pmap(api.pmap(api.vmap(functools.partial(fun1, do_print=True))))
    res = pmap_fun1(matrix)
    hcb.barrier_wait()
    expected_res = api.pmap(api.pmap(api.vmap(functools.partial(fun1, do_print=False))))(matrix)
    self.assertAllClose(expected_res, res, check_dtypes=False)
    # Assertion text is for 2 devices (also works for 1 device)
    assertMultiDeviceOutputEqual(self, """
        device: cpu:0 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
        [0.00 2.00 4.00]
        device: cpu:1 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
        [200.00 202.00 204.00]""")
    testing_stream.reset()

  @ignore_jit_of_pmap_warning()
  def test_pmap_pmap_extra(self):
    """pmap of a pmap surrounded by extra code."""
    # A matrix M[ij] = i * 10 + j
    nr_devices = api.local_device_count()
    if nr_devices != 2:
      raise SkipTest("test works only on 2 devices")
    shape = (2, 1, 3)
    matrix = np.fromfunction(lambda i, j, k: 100. * i + 10. * j + k, shape,
                             dtype=np.float32)
    def fun(xv, do_print=False):
      # This will be printed on all devices, with shape [1, 3]
      xv = maybe_print(do_print, xv + 1., "before", tap_with_device=True)
      res = api.pmap(lambda x: maybe_print(do_print, x * 2., "inside", tap_with_device=True))(xv)
      # This will be printed on all devices, with shape [1, 3]
      return maybe_print(do_print, res + 1., "after", tap_with_device=True)

    res = api.pmap(functools.partial(fun, do_print=True))(matrix)
    self.assertAllClose(fun(matrix, do_print=False), res, check_dtypes=False)
    hcb.barrier_wait()
    # Assertion text is for 2 devices (also works for 1 device)
    assertMultiDeviceOutputEqual(self, """
      device: cpu:0 what: before
      [[1.00 2.00 3.00]]
      device: cpu:0 what: inside
      [2.00 4.00 6.00]
      device: cpu:0 what: after
      [[3.00 5.00 7.00]]
      device: cpu:1 what: before
      [[101.00 102.00 103.00]]
      device: cpu:1 what: inside
      [202.00 204.00 206.00]
      device: cpu:1 what: after
      [[203.00 205.00 207.00]]""")

    testing_stream.reset()

  def test_jvp_pmap_vmap(self):
    # A matrix M[ijk] = i * 100 + j * 10 * k
    nr_devices = api.local_device_count()
    shape = (nr_devices, 2, 3)
    matrix = np.fromfunction(lambda i, j, k: 100. * i + 10. * j + k, shape,
                             dtype=np.float32)
    def fun(xv, do_print=False):
      # x: f32[3]
      return api.jvp(api.pmap(api.vmap(lambda x: maybe_print(do_print, x * 2., "x * 2", tap_with_device=True))),
                     (xv,), (.1 * jnp.ones_like(xv),))

    res = fun(matrix, do_print=True)
    hcb.barrier_wait()
    expected_res = fun(matrix, do_print=False)
    self.assertAllClose(expected_res, res, check_dtypes=False)
    # Assertion text is for 2 devices (also works for 1 device)
    # Device 0 will get to execute api.jvp(api.vmap(...)) for matrix[0, :, :]
    assertMultiDeviceOutputEqual(self, """
      device: cpu:0 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
      [[ 0.00  2.00  4.00]
       [20.00 22.00 24.00]]
      device: cpu:0 transforms: [('batch', {'batch_dims': (0,)}), 'jvp'] what: x * 2
      [[0.20 0.20 0.20]
       [0.20 0.20 0.20]]
      device: cpu:1 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
      [[200.00 202.00 204.00]
       [220.00 222.00 224.00]]
      device: cpu:1 transforms: [('batch', {'batch_dims': (0,)}), 'jvp'] what: x * 2
      [[0.20 0.20 0.20]
       [0.20 0.20 0.20]]""")
    testing_stream.reset()

  def test_vmap_pmap(self):
    # A matrix M[ijk] = i * 100 + j * 10 * k
    nr_devices = api.local_device_count()
    shape = (2, nr_devices, 3)
    matrix = np.fromfunction(lambda i, j, k: 100. * i + 10. * j + k, shape,
                             dtype=np.float32)
    def fun(xv, do_print=False):
      # x: f32[3]
      return api.vmap(api.pmap(lambda x: maybe_print(do_print, x * 2., "x * 2", tap_with_device=True)))(xv)

    res = fun(matrix, do_print=True)
    hcb.barrier_wait()
    expected_res = fun(matrix, do_print=False)
    self.assertAllClose(expected_res, res, check_dtypes=False)
    # Assertion text is for 2 devices (also works for 1 device)
    # Device 0 will get to execute api.jvp(api.vmap(...)) for matrix[:, 0, :]
    assertMultiDeviceOutputEqual(self, """
      device: cpu:0 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
      [[  0.00   2.00   4.00]
       [200.00 202.00 204.00]]
      device: cpu:1 transforms: [('batch', {'batch_dims': (0,)})] what: x * 2
      [[ 20.00  22.00  24.00]
       [220.00 222.00 224.00]]""")
    testing_stream.reset()

  @ignore_jit_of_pmap_warning()
  def test_jit_pmap_extra(self):
    """jit of a pmap surrounded by extra code."""
    # A matrix M[ij] = i * 10 + j
    nr_devices = api.local_device_count()
    shape = (nr_devices, 3)
    matrix = np.fromfunction(lambda i, j: 10. * i + j, shape,
                             dtype=np.float32)

    def fun(xv, do_print=False):
      # This will be printed on all devices with shape (nr_devices, 3)
      xv = maybe_print(do_print, xv + 1., "before", tap_with_device=True)
      res = api.pmap(lambda x: maybe_print(do_print, x * 2., "inside", tap_with_device=True))(xv)
      # This will be printed on all devices with shape (nr_devices, 3)
      return maybe_print(do_print, res + 1., "after", tap_with_device=True)

    res = api.jit(functools.partial(fun, do_print=True))(matrix)
    self.assertAllClose(fun(matrix, do_print=False), res, check_dtypes=False)
    hcb.barrier_wait()
    if api.device_count() == 2:
      assertMultiDeviceOutputEqual(self, """
        device: cpu:0 what: before
        [[ 1.00  2.00  3.00]
         [11.00 12.00 13.00]]
        device: cpu:0 what: inside
        [2.00 4.00 6.00]
        device: cpu:0 what: after
        [[ 3.00  5.00  7.00]
         [23.00 25.00 27.00]]
        device: cpu:1 what: before
        [[ 1.00  2.00  3.00]
         [11.00 12.00 13.00]]
        device: cpu:1 what: inside
        [22.00 24.00 26.00]
        device: cpu:1 what: after
        [[ 3.00  5.00  7.00]
         [23.00 25.00 27.00]]""")
    else:
      assert api.device_count() == 1
      assertMultiDeviceOutputEqual(self, """
        device: cpu:0 what: before
        [[1.00 2.00 3.00]]
        device: cpu:0 what: inside
        [2.00 4.00 6.00]
        device: cpu:0 what: after
        [[3.00 5.00 7.00]]""")

    testing_stream.reset()

  def test_cond_pmap(self):
    raise SkipTest("cond of pmap does not work in JAX. Issue #5178.")
    # A matrix M[ij] = i * 10 + j
    nr_devices = api.local_device_count()
    shape = (nr_devices, 3)
    matrix = np.fromfunction(lambda i, j: 10. * i + j, shape,
                             dtype=np.float32)
    def fun1(x, do_print=False):
      return maybe_print(do_print, x * 2., "x * 2")
    def fun2(cond, xv, do_print=False):
      return lax.cond(cond, api.pmap(functools.partial(fun1, do_print=do_print)),
                      lambda xv: xv, xv)
    res = fun2(True, matrix)
    self.assertAllClose(fun2(True, matrix, do_print=False), res, check_dtypes=False)
    hcb.barrier_wait()
    assertMultiLineStrippedEqual(self, """
        TBD""", testing_stream.output)
    testing_stream.reset()

  def test_scan_custom_jvp(self):
    """custom JVP, inside scan.
    This exercises the custom_jvp_call_jaxpr primitives."""
    @api.custom_jvp
    def f(x):
      return x * hcb.id_print(x, output_stream=testing_stream, what="x")

    @f.defjvp
    def f_jvp(primals, tangents):
      x, = primals
      x_dot, = tangents
      primal_out = f(x)
      tangent_out = 3. * x * hcb.id_print(x_dot, output_stream=testing_stream, what="x_dot")
      return primal_out, tangent_out

    def g(x):
      # Sum f(x_i)
      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
                      x)[0]

    arg = np.full((2,), 0.7)
    self.assertAllClose(0.7 * 0.7 * 2, g(arg))
    hcb.barrier_wait()
    self.assertMultiLineStrippedEqual("""
        what: x
        0.7
        what: x
        0.7""", testing_stream.output)
    testing_stream.reset()

    self.assertAllClose(np.array([2.1, 2.1]), api.grad(g)(arg), check_dtypes=False)
    hcb.barrier_wait()
    self.assertMultiLineStrippedEqual("""
        what: x
        0.7
        what: x
        0.7
        transforms: ['transpose'] what: x_dot
        2.1
        transforms: ['transpose'] what: x_dot
        2.1""", testing_stream.output)

  def test_scan_custom_vjp(self):
    """custom VJP, inside scan.
    This exercises the custom_vjp_call_jaxpr primitives."""
    @api.custom_vjp
    def f(x):
      return x * hcb.id_print(x, output_stream=testing_stream, what="x")

    # f_fwd: a -> (b, residual)
    def f_fwd(x):
      return f(x), 3. * x
    # f_bwd: (residual, CT b) -> [CT a]
    def f_bwd(residual, ct_b):
      return residual * hcb.id_print(ct_b, output_stream=testing_stream, what="ct_b"),

    f.defvjp(f_fwd, f_bwd)

    def g(x):
      # Sum f(x_i)
      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
                      x)[0]

    arg = np.full((2,), 0.7)

    self.assertAllClose(0.7 * 0.7 * 2, g(arg))
    hcb.barrier_wait()
    self.assertMultiLineStrippedEqual("""
        what: x
        0.7
        what: x
        0.7""", testing_stream.output)
    testing_stream.reset()

    self.assertAllClose(np.array([2.1, 2.1]), api.grad(g)(arg), check_dtypes=False)
    hcb.barrier_wait()
    self.assertMultiLineStrippedEqual("""
        what: x
        0.7
        what: x
        0.7
        what: ct_b
        1.
        what: ct_b
        1.""", testing_stream.output)

  def test_mask(self):
    # TODO(necula)
    raise SkipTest("masking has regressed")
    @functools.partial(api.mask, in_shapes=['n'], out_shape='')
    def padded_sum(x):
      return jnp.sum(hcb.id_print(x, what="x", output_stream=testing_stream))
    args = [jnp.arange(4)], dict(n=np.int64(2))
    assertMultiLineStrippedEqual(self, """
        { lambda c f ; a b.
          let d = lt c b
              e = id_tap[ func=_print
                          logical_shapes=[(Traced<ShapedArray(int32[]):JaxprTrace(level=0/0)>,)]
                          transforms=('mask',)
                          what=x ] a
              g = select d e f
              h = reduce_sum[ axes=(0,) ] g
          in (h,) }""", str(api.make_jaxpr(padded_sum)(*args)))

    _ = padded_sum(*args)
    self.assertMultiLineStrippedEqual("""
        logical_shapes: [(2,)] transforms: ['mask',) what: x
        [0 1 2 3]
          """, testing_stream.output)
    testing_stream.reset()

  def test_callback_delay(self):
    hcb.callback_extra = lambda dev: time.sleep(1)

    def func(x):
      for i in range(5):
        x = hcb.id_print(x * i, what="x times i")
      return x

    api.jit(func)(np.arange(6, dtype=np.float32).reshape((2, 3)))

  def test_callback_delay_barrier(self):
    hcb.callback_extra = lambda dev: time.sleep(2)

    def func(x):
      for i in range(1, 4):
        x = hcb.id_print(x * i, what="x times i", output_stream=testing_stream)
      return x

    api.jit(func)(np.arange(6, dtype=np.float32).reshape((2, 3)))
    # Wait for the results
    hcb.barrier_wait()
    expected = """
        what: x times i
        [[0. 1. 2.]
        [3. 4. 5.]]
        what: x times i
        [[ 0.  2.  4.]
        [ 6.  8. 10.]]
        what: x times i
        [[ 0.  6. 12.]
        [18. 24. 30.]]"""
    self.assertMultiLineStrippedEqual(expected, testing_stream.output)
    testing_stream.reset()
    # Call again
    api.jit(func)(np.arange(6, dtype=np.float32).reshape((2, 3)))
    hcb.barrier_wait()
    self.assertMultiLineStrippedEqual(expected, testing_stream.output)


  def test_error_bad_consumer_id(self):
    """Try to use reserved consumer ID 0.

    Check that we get the proper error from the runtime."""
    comp = xla_bridge.make_computation_builder(self._testMethodName)
    token = hcb.xops.CreateToken(comp)
    hcb._initialize_outfeed_receiver()  # Needed if this is the sole test
    with self.assertRaisesRegex(RuntimeError,
                                "Consumer ID cannot be a reserved value: 0"):
      hcb._outfeed_receiver.receiver.add_outfeed(
          comp, token, 0,
          [xla_bridge.constant(comp, np.zeros((2, 3), dtype=np.float32))])

  def test_error_different_shapes(self):
    """Try to register different shapes for the same consumer ID."""
    comp = xla_bridge.make_computation_builder(self._testMethodName)
    token = hcb.xops.CreateToken(comp)
    hcb._initialize_outfeed_receiver()  # Needed if this is the sole test
    hcb._outfeed_receiver.receiver.add_outfeed(
        comp, token, 123,
        [xla_bridge.constant(comp, np.zeros((2, 3), dtype=np.float32))])
    with self.assertRaisesRegex(
        RuntimeError, ".*does not match previous shape element_type.*"):
      hcb._outfeed_receiver.receiver.add_outfeed(
          comp, token, 123,
          [xla_bridge.constant(comp, np.zeros((2, 3), dtype=np.int32))])
    with self.assertRaisesRegex(
        RuntimeError, ".*does not match previous shape element_type.*"):
      hcb._outfeed_receiver.receiver.add_outfeed(
          comp, token, 123,
          [xla_bridge.constant(comp, np.zeros((2,), dtype=np.float32))])

  def test_id_tap_deprecated_kwargs(self):
    def func(x, transforms, y):
      pass
    with self.assertWarnsRegex(
        FutureWarning, r"Support for \*\*kwargs in ``id_tap``"):
      hcb.id_tap(func, 1, y=2)

  def test_odeint(self):
    # TODO: find a smaller repro for bug #4015
    # Seems to be xla_call(scan(xla_call)), all under grad.
    from jax.experimental.ode import odeint

    def f(x, t, k):
      x = hcb.id_print(x)
      return -k * x

    def loss(k=1.0):
      t = jnp.linspace(0, 0.001, num=2)
      xs = odeint(f, 1.0, t, k)
      return xs[-1]

    api.grad(loss)(1.0)  # should not fail

  def test_remat(self):
    def f(i, k):
      x = hcb.id_print(k + i, output_stream=testing_stream)
      return k * x

    def loss(k):
      return lax.fori_loop(0, 2, api.remat(f), k)
    print(loss(3))
    hcb.barrier_wait()
    expected = """
      3
      10"""
    self.assertMultiLineStrippedEqual(expected, testing_stream.output)


class OutfeedRewriterTest(jtu.JaxTestCase):

  def assertRewrite(self, expected: str, func: Callable, args: Sequence,
                    has_input_token=True, has_output_token=True):
    """Check that the rewrite of func(*args) matches expected."""
    jaxpr = api.make_jaxpr(func)(*args)
    rewritten = hcb._rewrite_closed_jaxpr(jaxpr,  # noqa: F841
                                          has_input_token, has_output_token)
    # Since it is somewhat annoying to update the Jaxpr assertions when we change
    # the Jaxpr printing, we do not check these by default. It is recommended that
    # before making changes to the code generation and Jaxpr rewriting, turn on
    # the checking, update the expected Jaxpr, and then make the changes.
    # assertMultiLineStrippedEqual(self, expected, str(rewritten))
    del jaxpr

  def test_no_outfeed(self):
    self.assertRewrite("""
        { lambda  ; a.
          let b = mul a a
              c = add a b
          in (c,) }""", lambda x: x + x * x, [0], has_input_token=False,
                              has_output_token=False)
    self.assertRewrite("""
        { lambda  ; a d.
          let b = mul a a
              c = add a b
          in (c,) }""", lambda x: x + x * x, [0], has_output_token=False)
    self.assertRewrite("""
        { lambda  ; a d.
          let b = mul a a
              c = add a b
          in (c, d) }""", lambda x: x + x * x, [0])

  def test_simple_outfeed(self):
    self.assertRewrite("""
        { lambda  ; a d.
          let b = add a a
              c e = id_tap[ arg_treedef_=*
                            has_token_=True
                            nr_tapped_args_=1
                            tap_func_=_print  ] b d
          in (c, e) }""", lambda x: hcb.id_print(x + x), [0])

  def test_simple_outfeed_without_input_token(self):
    self.assertRewrite("""
        { lambda  ; a b.
          let e = create_token a b
              c = add a b
              d f = id_tap[ arg_treedef_=*
                            has_token_=True
                            nr_tapped_args_=1
                            tap_func_=_print  ] c e
          in (d,) }""", lambda x1, x2: hcb.id_print(x1 + x2), [1, 2],
                        has_input_token=False, has_output_token=False)

  def test_simple_outfeed_without_input_token_nor_invars(self):
    self.assertRewrite("""
        { lambda  ; .
          let b = create_token
              a c = id_tap[ arg_treedef_=*
                            has_token_=True
                            nr_tapped_args_=1
                            tap_func_=_print  ] 42 b
          in (a,) }""", lambda: hcb.id_print(42), [],
                        has_input_token=False, has_output_token=False)

  def test_multiple_tap_without_dependencies(self):
    def f(x):
      hcb.id_print(x, what="x")
      hcb.id_print(x + 1, what="x + 1")
      return 2

    self.assertRewrite("""
        { lambda  ; a c.
          let _ d = id_tap[ arg_treedef_=*
                            has_token_=True
                            nr_tapped_args_=1
                            tap_func_=_print   what='x') ] a c
              b = add a 1
              _ e = id_tap[ arg_treedef_=*
                            has_token_=True
                            nr_tapped_args_=1
                            tap_func_=_print   what='x + 1') ] b d
          in (2, e) }""", f, [1])

  def test_cond(self):
    y = jnp.ones(5)  # captured const
    def func(x, z):
      return lax.cond(z > 0, (1, 2), lambda a: (a[0], jnp.zeros(5)),
                      z, lambda a: (hcb.id_print(a), y))
    self.assertRewrite("""
        { lambda a ; b c h.
          let d = gt c 0
              e = convert_element_type[ new_dtype=int32
                                        weak_type=False ] d
              f g i = cond[ branches=( { lambda  ; a b c d f.
                                         let e g = id_tap[ arg_treedef_=*
                                                           has_token_=True
                                                           nr_tapped_args_=1
                                                           tap_func_=_print  ] d f
                                         in (e, a, g) }
                                       { lambda  ; f_ a b c g.
                                         let d = broadcast_in_dim[ broadcast_dimensions=(  )
                                                                   shape=(5,) ] 0.00
                                         in (a, d, g) } )
                            linear=(False, False, False, False, False) ] e a 1 2 c h
          in (f, g, i) }""", func, [y, 5])

  def test_while(self):
    ct_body = jnp.ones(5, np.float32)  # captured const for the body
    ct_cond = jnp.ones(5, np.float32)  # captured const for the conditional

    def func(x):
      # x: f32[5]
      # c: (f32[5], f32)
      return lax.while_loop(lambda c: c[1] < jnp.sum(c[0] + ct_cond),
                            lambda c: (ct_body, hcb.id_print(c[1]) + 1.),
                            (x, np.float32(1.)))
    self.assertRewrite("""
        { lambda a b ; c f.
          let d e g = while[ body_jaxpr={ lambda  ; a b c f.
                                          let d g = id_tap[ arg_treedef_=*
                                                            has_token_=True
                                                            nr_tapped_args_=1
                                                            tap_func_=_print  ] c f
                                              e = add d 1.00
                                          in (a, e, g) }
                             body_nconsts=1
                             cond_jaxpr={ lambda  ; a b c g.
                                          let d = add b a
                                              e = reduce_sum[ axes=(0,) ] d
                                              f = lt c e
                                          in (f,) }
                             cond_nconsts=1 ] a b c 1.00 f
          in (d, e, g) }""", func, [ct_body])

  def test_while_pred_outfeed(self):
    """A while with outfeed in the pred."""
    ct_body = jnp.ones(5)  # captured const for the body
    ct_cond = jnp.ones(2) # captured const for the conditional

    def func(x):
      return lax.while_loop(lambda c: hcb.id_print(ct_cond, result=c[1]) < 5,
                            lambda c: (ct_body, hcb.id_print(c[1]) + 1),
                            (x, 1))

    self.assertRewrite("""
        { lambda a b ; c f.
          let h i = xla_call[ call_jaxpr={ lambda  ; a b c f.
                                           let _ d g = id_tap[ arg_treedef_=*
                                                               has_token_=True
                                                               nr_tapped_args_=1
                                                               tap_func_=_print  ] a c f
                                               e = lt d 5
                                           in (e, g) }
                              donated_invars=(False, False, False, False)
                              name=cond_before ] a c 1 f
              y d e g =
                while[ body_jaxpr={ lambda  ; n o p q r s.
                                    let t u v = xla_call[ call_jaxpr={ lambda  ; a b c f.
                                                                       let d g = id_tap[ arg_treedef_=*
                                                                                         has_token_=True
                                                                                         nr_tapped_args_=1
                                                                                         tap_func_=_print  ] c f
                                                                           e = add d 1
                                                                       in (a, e, g) }
                                                          donated_invars=(False, False, False, False)
                                                          name=body ] o q r s
                                        w x = xla_call[ call_jaxpr={ lambda  ; a b c f.
                                                                     let _ d g = id_tap[ arg_treedef_=*
                                                                                         has_token_=True
                                                                                         nr_tapped_args_=1
                                                                                         tap_func_=_print  ] a c f
                                                                         e = lt d 5
                                                                     in (e, g) }
                                                        donated_invars=(False, False, False, False)
                                                        name=cond_body ] n t u v
                                    in (w, t, u, x) }
                       body_nconsts=2
                       cond_jaxpr={ lambda  ; j k l m.
                                    let
                                    in (j,) }
                       cond_nconsts=0 ] a b h c 1 i
          in (d, e, g) }""", func, [ct_body])

  def test_scan(self):
    y = jnp.ones(5)  # captured const
    def func(x):
      return lax.scan(lambda c, a: (hcb.id_print(c), y), (1, 2), x)
    self.assertRewrite("""
        { lambda a ; b f.
          let c d g e =
                scan[ jaxpr={ lambda  ; a b c g d.
                              let e f h = id_tap[ arg_treedef_=PyTreeDef(tuple, [*,*])
                                                  has_token_=True
                                                  nr_tapped_args_=2
                                                  tap_func_=_print  ] b c g
                              in (e, f, h, a) }
                      length=5
                      linear=(False, False, False, False, False)
                      num_carry=3
                      num_consts=1
                      reverse=False
                      unroll=1 ] a 1 2 f b
          in (c, d, e, g) }""", func, [y])

  def test_scan_custom_jvp(self):
    """custom JVP, inside scan.
    This exercises the custom_jvp_call_jaxpr primitives."""
    @api.custom_jvp
    def f(x):
      return x * hcb.id_print(x)

    @f.defjvp
    def f_jvp(primals, tangents):
      x, = primals
      x_dot, = tangents
      primal_out = f(x)
      tangent_out = 3. * x * hcb.id_print(x_dot)
      return primal_out, tangent_out

    def g(x):
      # Sum f(x_i)
      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
                      x)[0]

    arg = np.full((5,), 0.7)
    self.assertRewrite("""
      { lambda  ; a c.
        let b d _ = scan[ jaxpr={ lambda  ; a e b.
                                  let c f = custom_jvp_call_jaxpr[ fun_jaxpr={ lambda  ; a d.
                                                                               let b e = id_tap[ arg_treedef_=*
                                                                                                 has_token_=True
                                                                                                 nr_tapped_args_=1
                                                                                                 tap_func_=_print  ] a d
                                                                                   c = mul a b
                                                                               in (c, e) }
                                                                   num_consts=0 ] b e
                                      d = add a c
                                  in (d, f, 0.00) }
                          length=5
                          linear=(False, False, False)
                          num_carry=2
                          num_consts=0
                          reverse=False
                          unroll=1 ] 0.00 c a
        in (b, d) }""", g, [arg])
    self.assertRewrite("""
      { lambda  ; a d.
        let _ _ e _ b =
              scan[ jaxpr={ lambda  ; a b h c d.
                            let e i = custom_jvp_call_jaxpr[ fun_jaxpr={ lambda  ; a d.
                                                                         let b e = id_tap[ arg_treedef_=*
                                                                                           has_token_=True
                                                                                           nr_tapped_args_=1
                                                                                           tap_func_=_print  ] a d
                                                                             c = mul a b
                                                                         in (c, e) }
                                                             num_consts=0 ] c h
                                f = add a e
                                g = mul c 3.00
                            in (f, *, i, 0.00, g) }
                    length=5
                    linear=(False, True, False, True, False)
                    num_carry=3
                    num_consts=0
                    reverse=False
                    unroll=1 ] 0.00 * d a *
            _ _ f _ c =
              scan[ jaxpr={ lambda  ; a b g c d.
                            let e = mul b d
                                f h = id_tap[ arg_treedef_=*
                                              has_token_=True
                                              nr_tapped_args_=1
                                              tap_func_=_print
                                              transforms=(('transpose',),) ] e g
                            in (*, b, h, *, f) }
                    length=5
                    linear=(True, True, True, False, False)
                    num_carry=3
                    num_consts=0
                    reverse=True
                    unroll=1 ] * 1.00 e * b
        in (c, f) }""", api.grad(g), [arg])

  def test_scan_custom_vjp(self):
    """custom VJP, inside scan.
    This exercises the custom_vjp_call_jaxpr primitives."""
    @api.custom_vjp
    def f(x):
      return x * hcb.id_print(x)

    # f_fwd: a -> (b, residual)
    def f_fwd(x):
      return f(x), 3. * x
    # f_bwd: (residual, CT b) -> [CT a]
    def f_bwd(residual, ct_b):
      return residual * hcb.id_print(ct_b),

    f.defvjp(f_fwd, f_bwd)

    def g(x):
      # Sum f(x_i)
      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
                      x)[0]

    arg = np.full((2,), 0.7)
    self.assertRewrite("""
      { lambda  ; a c.
        let b d _ = scan[ jaxpr={ lambda  ; a e b.
                                  let c f = custom_vjp_call_jaxpr[
                                                                   fun_jaxpr={ lambda  ; a d.
                                                                               let b e = id_tap[ arg_treedef_=*
                                                                                                 has_token_=True
                                                                                                 nr_tapped_args_=1
                                                                                                 tap_func_=_print  ] a d
                                                                                   c = mul a b
                                                                               in (c, e) }
                                                                   num_consts=0
                                                                   ] b e
                                      d = add a c
                                  in (d, f, 0.00) }
                          length=2
                          linear=(False, False, False)
                          num_carry=2
                          num_consts=0
                          reverse=False
                          unroll=1 ] 0.00 c a
        in (b, d) }""", g, [arg])
    self.assertRewrite("""
      { lambda  ; a d.
        let _ _ e _ b =
              scan[ jaxpr={ lambda  ; a b h c d.
                            let e i = custom_vjp_call_jaxpr[
                                                             fun_jaxpr={ lambda  ; a d.
                                                                         let b e = id_tap[ arg_treedef_=*
                                                                                           has_token_=True
                                                                                           nr_tapped_args_=1
                                                                                           tap_func_=_print  ] a d
                                                                             c = mul a b
                                                                         in (c, e) }
                                                             num_consts=0
                                                             ] c h
                                f = add a e
                                g = mul c 3.00
                            in (f, *, i, 0.00, g) }
                    length=2
                    linear=(False, True, False, True, False)
                    num_carry=3
                    num_consts=0
                    reverse=False
                    unroll=1 ] 0.00 * d a *
            _ _ f _ c =
              scan[ jaxpr={ lambda  ; a b g c d.
                            let e h = id_tap[ arg_treedef_=*
                                              has_token_=True
                                              nr_tapped_args_=1
                                              tap_func_=_print  ] b g
                                f = mul d e
                            in (*, b, h, *, f) }
                    length=2
                    linear=(True, True, True, False, False)
                    num_carry=3
                    num_consts=0
                    reverse=True
                    unroll=1 ] * 1.00 e * b
        in (c, f) }""", api.grad(g), [arg])

  def test_remat_loop(self):
    def f(k, x):
      x = hcb.id_print(k + x)
      return -k * x

    def loss(k):
      return lax.fori_loop(0, 1, api.remat(f), k)

    self.assertRewrite("""
        { lambda  ; a c.
          let _ _ b d =
                while[ body_jaxpr={ lambda  ; a b c f.
                                    let d = add a 1
                                        e g = remat_call[ call_jaxpr={ lambda  ; a b g.
                                                                       let c = add a b
                                                                           d h = id_tap[ arg_treedef_=*
                                                                                         has_token_=True
                                                                                         nr_tapped_args_=1
                                                                                         tap_func_=_print  ] c g
                                                                           e = neg a
                                                                           f = mul e d
                                                                       in (f, h) }
                                                          concrete=False
                                                          name=f ] a c f
                                    in (d, b, e, g) }
                       body_nconsts=0
                       cond_jaxpr={ lambda  ; a b c e.
                                    let d = lt a b
                                    in (d,) }
                       cond_nconsts=0 ] 0 1 a c
          in (b, d) }""", loss, [2])

  def test_pmap(self):
    def f(xv):
      api.pmap(lambda x: jnp.sin(hcb.id_print(x, tap_with_device=True)),
               axis_name="i")(xv)

    self.assertRewrite("""
      { lambda  ; a b.
        let _ c = xla_pmap[ axis_name=i
                            axis_size=1
                            backend=None
                            call_jaxpr={ lambda  ; a e.
                                         let b f = id_tap[ arg_treedef_=*
                                                           has_token_=True
                                                           nr_tapped_args_=1
                                                           tap_func_=_print
                                                           tap_with_device_=True ] a e
                                             c = convert_element_type[ new_dtype=float32
                                                                       weak_type=False ] b
                                             d = sin c
                                         in (d, f) }
                            devices=None
                            donated_invars=(False, False)
                            global_arg_shapes=(None,)
                            global_axis_size=None
                            in_axes=(0, 0)
                            name=<lambda>
                            out_axes=(0, 0) ] a b
        in (c,) }""", f, [np.array([2])])

if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())