code: def check_close(x, y, tol=1e-3):
      assert jnp.shape(x) == jnp.shape(y)
      # TODO(dougalm): re-enable once we've tackled the less pendantic bugs
      # assert x.dtype == y.dtype
      assert jnp.allclose(
          x, y, rtol=tol, atol=tol
      ), f"Value mismatch:\n{x}\n  vs\n{y}\n"
  
imports: from collections import namedtuple
  from functools import partial
  import numpy.random
  from absl.testing import absltest
  from absl.testing import parameterized
  import itertools
  import jax.numpy
  from jax import jit, jvp, vjp
  import jax._src.test_util
  from jax import config
  from jax._src.util import unzip2, safe_zip, safe_map
label_desc: Use image node_id 23 for calling a global function with example usage: check_close(x, y, tol) without return types