code: def test_convert(
      self,
      module_fn: ModuleFn,
      shape,
      dtype,
      init: bool,
      tf_transform,
      jax_transform,
  ):
      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):
          return module_fn()(x)
  
      f = hk.transform_with_state(g)
  
      atol = CUSTOM_ATOL.get(
          descriptors.module_type(module_fn), DEFAULT_ATOL
      )
      assert_allclose = functools.partial(
          np.testing.assert_allclose, atol=atol
      )
  
      get = lambda t: jax.tree_util.tree_map(lambda x: x.numpy(), t)
  
      if init:
          init_jax = jax_transform(f.init)
          init_tf = tf_transform(jax2tf.convert(f.init))
          jax.tree_util.tree_map(
              assert_allclose, init_jax(rng, x), get(init_tf(rng, x))
          )
  
      else:
          params, state = f.init(rng, x)
          apply_jax = jax_transform(f.apply)
          apply_tf = tf_transform(jax2tf.convert(f.apply))
          jax.tree_util.tree_map(
              assert_allclose,
              apply_jax(params, state, rng, x),
              get(apply_tf(params, state, rng, x)),
          )
  
file_path: dm-haiku/haiku/_src/integration/jax2tf_test.py
imports: import functools
  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
  from jax.experimental import jax2tf
  import jax.numpy
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling the JaxToTfTest obj's underlying member method code with example usage: obj.test_convert(module_fn, shape, dtype, init, tf_transform, jax_transform) and returns: module_fn