class_docstr: Test cases for AxialResNet.
class_docstr_tok: ["Test","cases","for","AxialResNet","."]
code: def test_self_attention_with_1d_relative_pos_unacceptable_num_features(
      self,
  ):
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 6, 5))
      # The channel size should be divisible by `2 * num_heads`, which is not:
      with self.assertRaises(ValueError):
          axial_resnet.SelfAttentionWith1DRelativePos(num_heads=8).init(
              rng, x
          )
  
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 2 for calling the AxialResNetTest obj's underlying member method code with example usage: obj.test_self_attention_with_1d_relative_pos_unacceptable_num_features() without return types