class_docstr: Test cases for AxialResNet.
class_docstr_tok: ["Test","cases","for","AxialResNet","."]
code: def test_self_attention_with_1d_relative_pos_output_shape(self):
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 6, 128))
      sa_module = axial_resnet.SelfAttentionWith1DRelativePos(
          num_heads=8
      )
      y, _ = sa_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 1 for calling the AxialResNetTest obj's underlying member method code with example usage: obj.test_self_attention_with_1d_relative_pos_output_shape() without return types