code: def tree_assert_allclose(a, b, *, atol=1e-6):
      jax.tree_util.tree_map(
          functools.partial(np.testing.assert_allclose, atol=atol), a, b
      )
  
file_path: dm-haiku/haiku/_src/integration/numpy_inputs_test.py
imports: import functools
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  from haiku._src import test_utils
  from haiku._src.integration import descriptors
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: tree_assert_allclose(a, b) without return types