class_docstr: Tests for modules in attention_layers.py.
class_docstr_tok: ["Tests","for","modules","in","attention_layers.py","."]
docstr: Test axial_dot_product_attention function.
docstr_tok: ["Test","axial_dot_product_attention","function","."]
code: def test_relative_attention_bias(self, nd_shape):
      """Test axial_dot_product_attention function."""
      num_heads = 2
      bias_layer = attention_layers.RelativeAttentionBias(
          num_heads=num_heads,
          nd_shape=nd_shape,
          initializer=nn.initializers.normal(),
      )
      rng = random.PRNGKey(0)
      variables = bias_layer.init(rng)
      bias = bias_layer.apply(variables)
  
      length = np.prod(nd_shape)
      self.assertEqual((num_heads, length, length), bias.shape)
  
      bias_nd = bias.reshape((num_heads,) + nd_shape + nd_shape)
      for i in range(len(nd_shape)):
          bias_crop = bias_nd
          for _ in range(i + 1, len(nd_shape)):
              # Crop until last dim is dim to be checked.
              bias_crop = bias_crop[:, :, ..., 0]
          for k in range(nd_shape[i] - 1):
              np.testing.assert_array_equal(
                  bias_crop[:, k, ..., :-1],
                  bias_crop[:, k + 1, ..., 1:],
              )
          bias_nd = bias_nd[:, 0]
  
      # Now plug this bias into multi-head attention.
      layer = attention_layers.MultiHeadAttention(num_heads=num_heads)
      input_shape = (4, length, num_heads * 2)
      inputs = jnp.array(np.random.normal(size=input_shape))
      variables = layer.init(rng, inputs, inputs)
      layer.apply(variables, inputs, inputs, attention_bias=bias)
  
file_path: scenic/scenic/model_lib/layers/tests/test_attention_layers.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 8 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_relative_attention_bias(nd_shape) without return types