code: def test_big_resnet(self):
      x = random.normal(random.key(0), (1, 8, 8, 8))
      y, variables = init(big_resnet)(random.key(1), x)
      self.assertEqual(y.shape, (1, 8, 8, 8))
      param_shapes = unfreeze(
          jax.tree_util.tree_map(jnp.shape, variables["params"])
      )
      batch_stats_shapes = unfreeze(
          jax.tree_util.tree_map(jnp.shape, variables["batch_stats"])
      )
      self.assertEqual(
          param_shapes,
          {
              "conv_1": {"kernel": (10, 5, 3, 3, 8, 8)},
              "conv_2": {"kernel": (10, 5, 3, 3, 8, 8)},
              "bn_1": {"scale": (10, 5, 8), "bias": (10, 5, 8)},
              "bn_2": {"scale": (10, 5, 8), "bias": (10, 5, 8)},
          },
      )
      self.assertEqual(
          batch_stats_shapes,
          {
              "bn_1": {"var": (10, 5, 8), "mean": (10, 5, 8)},
              "bn_2": {"var": (10, 5, 8), "mean": (10, 5, 8)},
          },
      )
  
file_path: flax/tests/core/design/core_big_resnets_test.py
imports: from functools import partial
  import jax
  import numpy
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, init, lift, nn, unfreeze
label_desc: Use image node_id 1 for calling the BigResnetTest obj's underlying member method code with example usage: obj.test_big_resnet() without return types