code: def test_resnet_50_apply_equivalence(self):
      mod = flax_module.Module.create(resnet.ResNet50, 10)
      rng = jax.random.PRNGKey(42)
      x = jnp.ones([1, 224, 224, 3])
      f_haiku = transform.transform_with_state(
          lambda x: resnet.ResNet50(10)(x, is_training=True)
      )
      variables = dict(mod.init(rng, x, is_training=True))
      # Haiku and Flax have very different RNG key implementations, so parameter
      # initialisation does not match when using `flax_module`. There is no
      # mechanism to make Flax initialization match Haiku initialisation.
      hk_params = utils.flatten_flax_to_haiku(variables["params"])
      hk_state = utils.flatten_flax_to_haiku(variables["state"])
  
      for _ in range(5):
          out_flax, state = mod.apply(
              variables, x, is_training=True, mutable=["state"]
          )
          out_hk, hk_state = f_haiku.apply(hk_params, hk_state, None, x)
          np.testing.assert_array_equal(out_flax, out_hk)
          variables = {**variables, **state}
          assert_matches(variables["state"], hk_state)
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from haiku._src import base
  from haiku._src import filtering
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.flax import flax_module
  from haiku._src.flax import utils
  from haiku._src.nets import resnet
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 6 for calling the FlaxModuleTest obj's underlying member method code with example usage: obj.test_resnet_50_apply_equivalence() without return types