class_docstr: Test cases for AxialResNet.
class_docstr_tok: ["Test","cases","for","AxialResNet","."]
docstr: Tests AxialSelfAttention module given different attention axis.
docstr_tok: ["Tests","AxialSelfAttention","module","given","different","attention","axis","."]
code: def test_axial_self_attention_output_shape(self, attention_axis):
      """Tests AxialSelfAttention module given different attention axis."""
      rng = random.PRNGKey(0)
      x = jnp.ones((6, 8, 10, 128))
      axial_attention_configs = ml_collections.ConfigDict(
          {"num_heads": 4}
      )
      if attention_axis not in [1, 2]:
          with self.assertRaises(ValueError):
              axial_resnet.AxialSelfAttention(
                  attention_axis=attention_axis,
                  axial_attention_configs=axial_attention_configs,
              ).init_with_output(rng, x)
      else:
          asa_module = axial_resnet.AxialSelfAttention(
              attention_axis=attention_axis,
              axial_attention_configs=axial_attention_configs,
          )
          y, _ = asa_module.init_with_output(rng, x)
          self.assertEqual(y.shape, x.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 3 for calling the AxialResNetTest obj's underlying member method code with example usage: obj.test_axial_self_attention_output_shape(attention_axis) without return types