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_axial_dot_product_attention_has_expected_shape(
      self, q_shape, k_shape
  ):
      """Test axial_dot_product_attention function."""
      v_shape = k_shape[:-1] + (64,)
      expected_output_shape = q_shape[:-1] + (v_shape[-1],)
  
      query = jnp.array(np.random.normal(size=q_shape))
      key = jnp.array(np.random.normal(size=k_shape))
      value = jnp.array(np.random.normal(size=v_shape))
      y = attention_layers.axial_dot_product_attention(
          query, key, value, deterministic=True
      )
      # Test outputs shape:
      self.assertEqual(y.shape, expected_output_shape)
  
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 7 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_axial_dot_product_attention_has_expected_shape(q_shape, k_shape) without return types