class_docstr: Tests for modules in attention_layers.py.
class_docstr_tok: ["Tests","for","modules","in","attention_layers.py","."]
docstr: Tests Add1DPositionEmbedding.
docstr_tok: ["Tests","Add1DPositionEmbedding","."]
code: def test_add_1d_positional_embedding(self, posemb_init):
      """Tests Add1DPositionEmbedding."""
      rng = random.PRNGKey(0)
      input_shape = (4, 16, 32)
      inputs = jnp.array(np.random.normal(size=input_shape))
  
      # Test output after adding positional embedding.
      layer = attention_layers.Add1DPositionEmbedding(
          posemb_init=posemb_init
      )
      variables = layer.init(rng, inputs)
      outputs = layer.apply(variables, inputs)
  
      # Test output shape.
      self.assertEqual(outputs.shape, input_shape)
  
      if posemb_init is not None:
          # Test added learned 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"]["pos_embedding"].shape,
              (1,) + input_shape[1:],
          )
  
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 5 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_add_1d_positional_embedding(posemb_init) without return types