docstr: Test computed shape is equal to the layer output's shape.
docstr_tok: ["Test","computed","shape","is","equal","to","the","layer","output","'s","shape","."]
code: def test_compute_output_shape(
      self,
      query_dims,
      value_dims,
      key_dims,
      num_query_heads,
      num_key_value_heads,
  ):
      """Test computed shape is equal to the layer output's shape."""
      layer = layers.GroupedQueryAttention(
          num_query_heads=num_query_heads,
          num_key_value_heads=num_key_value_heads,
          head_dim=2,
      )
      batch_size = 7
      query_shape = (batch_size,) + query_dims
      value_shape = (batch_size,) + value_dims
      key_shape = (batch_size,) + key_dims if key_dims else None
  
      query = np.ones(query_shape)
      value = np.ones(value_shape)
      key = np.ones(key_shape) if key_shape else None
      output = layer(query=query, value=value, key=key)
      comp_output_shape = layer.compute_output_shape(
          query_shape, value_shape, key_shape
      )
      self.assertEqual(output.shape, comp_output_shape)
  
file_path: keras/keras/layers/attention/grouped_query_attention_test.py
imports: import numpy
  import pytest
  from absl.testing import parameterized
  from keras import backend
  from keras import initializers
  from keras import layers
  from keras import testing
label_desc: Use image node_id 2 for calling the GroupedQueryAttentionTest obj's underlying member method code with example usage: obj.test_compute_output_shape(query_dims, value_dims, key_dims, num_query_heads, num_key_value_heads) without return types