class_docstr: Tests for modules in attention_layers.py.
class_docstr_tok: ["Tests","for","modules","in","attention_layers.py","."]
docstr: Tests Add2DPositionEmbedding.
docstr_tok: ["Tests","Add2DPositionEmbedding","."]
code: def test_add_2d_positional_embedding(self):
      """Tests Add2DPositionEmbedding."""
      rng = random.PRNGKey(0)
      input_shape = (4, 8, 16, 32)
      inputs = jnp.ones(input_shape)
  
      # Test output after adding positional embedding.
      layer = attention_layers.Add2DPositionEmbedding(
          posemb_init=nn.initializers.ones
      )
      variables = layer.init(rng, inputs)
      outputs = layer.apply(variables, inputs)
  
      # Test output shape.
      self.assertEqual(outputs.shape, input_shape)
  
      # Test added embeddings.
      # Note that we initialize them with nn.initializers.ones.
      expected_added_pos_emb = jnp.ones(input_shape, dtype=inputs.dtype)
      added_pos_emb = outputs - inputs
      np.testing.assert_allclose(
          added_pos_emb, expected_added_pos_emb, atol=1e-6
      )
  
      # Test embeddings shape.
      self.assertEqual(
          variables["params"]["row_pos_embedding"].shape,
          (input_shape[2], input_shape[-1] // 2),
      )
      self.assertEqual(
          variables["params"]["col_pos_embedding"].shape,
          (input_shape[1], input_shape[-1] // 2),
      )
  
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 6 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_add_2d_positional_embedding() without return types