code: def test_hk_remat(self, module_fn: ModuleFn, shape, dtype):
      rng = jax.random.PRNGKey(42)
      if jnp.issubdtype(dtype, jnp.integer):
          x = jax.random.randint(rng, shape, 0, np.prod(shape), dtype)
      else:
          x = jax.random.uniform(rng, shape, dtype)
  
      def g(x, remat=False):
          mod = module_fn()
          if remat:
              mod = hk.remat(mod)
          out = mod(x)
          if isinstance(out, dict):
              out = out["loss"]
          return jnp.mean(out)
  
      f = hk.transform_with_state(g)
  
      assert_allclose = functools.partial(
          np.testing.assert_allclose, atol=1e-5
      )
  
      grad_jax_remat = jax.grad(jax.remat(f.apply), has_aux=True)
      grad_hk_remat = jax.grad(
          functools.partial(f.apply, remat=True), has_aux=True
      )
  
      params, state = f.init(rng, x)
      jax.tree_util.tree_map(
          assert_allclose,
          grad_jax_remat(params, state, rng, x),
          grad_hk_remat(params, state, rng, x),
      )
  
file_path: dm-haiku/haiku/_src/integration/hk_transforms_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  from haiku._src import stateful
  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 3 for calling the HaikuTransformsTest obj's underlying member method code with example usage: obj.test_hk_remat(module_fn, shape, dtype) and returns: jnp