code: def test_numpy_and_jax_results_close(
      self,
      module_fn: ModuleFn,
      shape: tuple[int, ...],
      dtype: jnp.dtype,
      np_params: bool,
      np_inputs: bool,
      close_over_params: bool,
  ):
      if not (np_params or np_inputs):
          self.skipTest("Pure JAX variants tested elsewhere")
  
      # TODO(b/257921991): Fix the timeouts here.
      if close_over_params and "UNITTEST_ON_FORGE" in os.environ:
          module_cls = get_module_cls(module_fn)
          if module_cls in (
              hk.nets.ResNet,
              hk.Conv2DTranspose,
              hk.LayerNorm,
              hk.Conv2DLSTM,
              hk.IdentityCore,
          ):
              self.skipTest(
                  "Close over tests for these modules take >5 minutes"
              )
  
      f = hk.transform_with_state(
          lambda x: module_fn()(x)
      )  # pylint: disable=unnecessary-lambda
  
      rng = jax.random.PRNGKey(42)
      x = jnp.ones(shape, dtype)
      params, state = f.init(rng, x)
      if close_over_params:
          apply_fn = functools.partial(f.apply, params, state)
          out, new_state = jax.jit(apply_fn)(rng, x)
      else:
          out, new_state = jax.jit(f.apply)(params, state, rng, x)
  
      if np_inputs:
          rng, x = jax.device_get((rng, x))
  
          with self.subTest("init"):
              params2, state2 = f.init(rng, x)
              tree_assert_allclose(params, params2)
              tree_assert_allclose(state, state2)
  
      with self.subTest("apply"):
          if np_params:
              params, state = jax.device_get((params, state))
  
          if close_over_params:
              apply_fn = functools.partial(f.apply, params, state)
              out2, new_state2 = jax.jit(apply_fn)(rng, x)
          else:
              out2, new_state2 = jax.jit(f.apply)(params, state, rng, x)
  
          tree_assert_allclose(out, out2)
          tree_assert_allclose(new_state, new_state2)
  
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 the NumpyInputsTest obj's underlying member method code with example usage: obj.test_numpy_and_jax_results_close(module_fn, shape, dtype, np_params, np_inputs, close_over_params) without return types