code: def test_local_stats(self, resnet_v2, bottleneck):
      def forward_fn(image):
          model = resnet.ResNet(
              [1, 1, 1, 1],
              10,
              resnet_v2=resnet_v2,
              bottleneck=bottleneck,
          )
          return model(image, is_training=False, test_local_stats=True)
  
      forward = transform.transform(forward_fn)
      rng = jax.random.PRNGKey(42)
      image = jnp.ones([2, 64, 64, 3])
      params = forward.init(rng, image)
      logits = forward.apply(params, None, image)
      self.assertEqual(logits.shape, (2, 10))
  
file_path: dm-haiku/haiku/_src/nets/resnet_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  from haiku._src import test_utils
  from haiku._src import transform
  from haiku._src.nets import resnet
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the ResnetTest obj's underlying member method code with example usage: obj.test_local_stats(resnet_v2, bottleneck) and returns: model