class_docstr: Test cases for AxialResNet.
class_docstr_tok: ["Test","cases","for","AxialResNet","."]
docstr: Tests AxialResNetStage module given different strides.
docstr_tok: ["Tests","AxialResNetStage","module","given","different","strides","."]
code: def test_axial_residual_stage_output_shape(
      self, strides, bottleneck, block_size, expected_output_shape
  ):
      """Tests AxialResNetStage module given different strides."""
      rng = random.PRNGKey(0)
      x = jnp.ones((10, 32, 32, 64))
      axial_attention_configs = ml_collections.ConfigDict(
          {"num_heads": 4}
      )
      aru_module = axial_resnet.AxialResNetStage(
          block_size=block_size,
          nout=128,
          first_stride=strides,
          bottleneck=bottleneck,
          axial_attention_configs=axial_attention_configs,
      )
      y, _ = aru_module.init_with_output(rng, x)
      self.assertEqual(y.shape, expected_output_shape)
  
file_path: scenic/scenic/projects/baselines/tests/test_axial_resnet.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from jax import random
  import jax.numpy
  import ml_collections
  from scenic.projects.baselines import axial_resnet
label_desc: Use image node_id 5 for calling the AxialResNetTest obj's underlying member method code with example usage: obj.test_axial_residual_stage_output_shape(strides, bottleneck, block_size, expected_output_shape) without return types