code: def test_resnet_50_init_equivalence_to_flax(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)
      )
      # We check that init is equivalent when passed the RNG used by Flax. There
      # is no mechanism currently to test the inverse (init the Flax module to
      # match Haiku init) because of how Flax and Haiku disagree on RNG key
      # splitting.
      hk_params, hk_state = f_haiku.init(flax_init_rng(rng), x)
      variables = dict(mod.init(rng, x, is_training=True))
      assert_matches(variables["params"], hk_params)
      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 5 for calling the FlaxModuleTest obj's underlying member method code with example usage: obj.test_resnet_50_init_equivalence_to_flax() without return types