code: def test_resnet(self):
      block_sizes = (2, 2)
      x = random.normal(random.key(0), (1, 64, 64, 3))
      y, variables = init(resnet)(
          random.key(1), x, block_sizes=block_sizes, features=16
      )
      param_shapes = unfreeze(
          jax.tree_util.tree_map(jnp.shape, variables["params"])
      )
      self.assertEqual(y.shape, (1, 1000))
  
      self.assertEqual(
          param_shapes,
          {
              "init_conv": {"kernel": (7, 7, 3, 16)},
              "init_bn": {"bias": (16,), "scale": (16,)},
              "out": {"kernel": (128, 1000), "bias": (1000,)},
              "block_0_0": {
                  "conv_1": {"kernel": (1, 1, 16, 16)},
                  "conv_2": {"kernel": (3, 3, 16, 64)},
                  "conv_3": {"kernel": (1, 1, 64, 64)},
                  "bn_1": {"bias": (16,), "scale": (16,)},
                  "bn_2": {"bias": (64,), "scale": (64,)},
                  "bn_3": {"bias": (64,), "scale": (64,)},
                  "proj_conv": {"kernel": (1, 1, 16, 64)},
                  "proj_bn": {"bias": (64,), "scale": (64,)},
              },
              "block_0_1": {
                  "conv_1": {"kernel": (1, 1, 64, 16)},
                  "conv_2": {"kernel": (3, 3, 16, 64)},
                  "conv_3": {"kernel": (1, 1, 64, 64)},
                  "bn_1": {"bias": (16,), "scale": (16,)},
                  "bn_2": {"bias": (64,), "scale": (64,)},
                  "bn_3": {"bias": (64,), "scale": (64,)},
              },
              "block_1_0": {
                  "conv_1": {"kernel": (1, 1, 64, 32)},
                  "conv_2": {"kernel": (3, 3, 32, 128)},
                  "conv_3": {"kernel": (1, 1, 128, 128)},
                  "bn_1": {"bias": (32,), "scale": (32,)},
                  "bn_2": {"bias": (128,), "scale": (128,)},
                  "bn_3": {"bias": (128,), "scale": (128,)},
                  "proj_conv": {"kernel": (1, 1, 64, 128)},
                  "proj_bn": {"bias": (128,), "scale": (128,)},
              },
              "block_1_1": {
                  "conv_1": {"kernel": (1, 1, 128, 32)},
                  "conv_2": {"kernel": (3, 3, 32, 128)},
                  "conv_3": {"kernel": (1, 1, 128, 128)},
                  "bn_1": {"bias": (32,), "scale": (32,)},
                  "bn_2": {"bias": (128,), "scale": (128,)},
                  "bn_3": {"bias": (128,), "scale": (128,)},
              },
          },
      )
  
file_path: flax/tests/core/design/core_resnet_test.py
imports: from functools import partial
  import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, init, nn, unfreeze
label_desc: Use image node_id 1 for calling the ResNetTest obj's underlying member method code with example usage: obj.test_resnet() without return types