Merge pull request #3706 from AI-Hypercomputer:add-embeddings-tests

PiperOrigin-RevId: 902870201

Author: Google-ML-Automation

tests/unit/embeddings_test.py

@@ -0,0 +1,202 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for embeddings.py."""
+
+import sys
+import unittest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxtext.layers import embeddings
+from maxtext.configs import pyconfig
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+
+class EmbedTest(unittest.TestCase):
+  """Tests for Embed."""
+
+  def setUp(self):
+    super().setUp()
+    self.rngs = nnx.Rngs(params=0)
+
+    config_arguments = {
+        "per_device_batch_size": 1.0,
+        "run_name": "test",
+        "enable_checkpointing": False,
+        "max_target_length": 128,
+    }
+    argv = [sys.argv[0], get_test_config_path()]
+    self.cfg = pyconfig.initialize(argv, **config_arguments)
+
+    devices_array = maxtext_utils.create_device_mesh(self.cfg)
+    self.mesh = jax.sharding.Mesh(devices_array, self.cfg.mesh_axes)
+
+  def test_basic_call(self):
+    num_embeddings = 100
+    num_features = 16
+    batch_size = 2
+    seq_len = 3
+
+    layer = embeddings.Embed(
+        num_embeddings=num_embeddings,
+        num_features=num_features,
+        config=self.cfg,
+        mesh=self.mesh,
+        rngs=self.rngs,
+    )
+
+    inputs = jnp.zeros((batch_size, seq_len), dtype=jnp.int32)
+    outputs = layer(inputs)
+
+    self.assertEqual(outputs.shape, (batch_size, seq_len, num_features))
+
+  def test_attend(self):
+    num_embeddings = 100
+    num_features = 16
+    batch_size = 2
+    seq_len = 3
+
+    layer = embeddings.Embed(
+        num_embeddings=num_embeddings,
+        num_features=num_features,
+        config=self.cfg,
+        mesh=self.mesh,
+        rngs=self.rngs,
+    )
+
+    query = jnp.ones((batch_size, seq_len, num_features))
+    outputs = layer.attend(query)
+
+    self.assertEqual(outputs.shape, (batch_size, seq_len, num_embeddings))
+
+
+class RotaryEmbeddingTest(unittest.TestCase):
+  """Tests for RotaryEmbedding."""
+
+  def setUp(self):
+    super().setUp()
+    self.rngs = nnx.Rngs(params=0)
+
+    config_arguments = {
+        "per_device_batch_size": 1.0,
+        "run_name": "test",
+        "enable_checkpointing": False,
+        "max_target_length": 128,
+    }
+    argv = [sys.argv[0], get_test_config_path()]
+    self.cfg = pyconfig.initialize(argv, **config_arguments)
+
+    devices_array = maxtext_utils.create_device_mesh(self.cfg)
+    self.mesh = jax.sharding.Mesh(devices_array, self.cfg.mesh_axes)
+
+  def test_basic_call(self):
+    layer = embeddings.RotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=4,
+        rngs=self.rngs,
+    )
+
+    inputs = jnp.ones((1, 2, 1, 4))
+    position = jnp.array([[0, 1]])
+
+    outputs = layer(inputs, position=position)
+
+    self.assertEqual(outputs.shape, (1, 2, 1, 4))
+
+    # Snapshot verification
+    expected = jnp.array([[[[1.0, 1.0, 1.0, 1.0]], [[-0.300781, 0.988281, 1.38281, 1.00781]]]])
+    np.testing.assert_allclose(outputs, expected, atol=1e-5)
+
+
+class LLaMARotaryEmbeddingTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.rngs = nnx.Rngs(params=0)
+
+    config_arguments = {
+        "per_device_batch_size": 1.0,
+        "run_name": "test",
+        "enable_checkpointing": False,
+        "max_target_length": 128,
+    }
+    argv = [sys.argv[0], get_test_config_path()]
+    self.cfg = pyconfig.initialize(argv, **config_arguments)
+
+    devices_array = maxtext_utils.create_device_mesh(self.cfg)
+    self.mesh = jax.sharding.Mesh(devices_array, self.cfg.mesh_axes)
+
+  def test_basic_call(self):
+    layer = embeddings.LLaMARotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=4,
+        use_scale=True,
+        rngs=self.rngs,
+    )
+    inputs = jnp.ones((1, 2, 1, 4))
+    position = jnp.array([[0, 1]])
+    outputs = layer(inputs, position=position)
+    self.assertEqual(outputs.shape, (1, 2, 1, 4))
+
+    # Snapshot verification
+    expected = jnp.array([[[[1.0, 1.0, 1.0, 1.0]], [[-0.300781, 1.38281, 0.988281, 1.00781]]]])
+    np.testing.assert_allclose(outputs, expected, atol=1e-5)
+
+
+class YarnRotaryEmbeddingTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.rngs = nnx.Rngs(params=0)
+
+    config_arguments = {
+        "per_device_batch_size": 1.0,
+        "run_name": "test",
+        "enable_checkpointing": False,
+        "max_target_length": 128,
+    }
+    argv = [sys.argv[0], get_test_config_path()]
+    self.cfg = pyconfig.initialize(argv, **config_arguments)
+
+    devices_array = maxtext_utils.create_device_mesh(self.cfg)
+    self.mesh = jax.sharding.Mesh(devices_array, self.cfg.mesh_axes)
+
+  def test_basic_call(self):
+    layer = embeddings.YarnRotaryEmbedding(
+        embedding_dims=4,
+        mesh=self.mesh,
+        max_position_embeddings=16384,
+        original_max_position_embeddings=4096,
+        rngs=self.rngs,
+    )
+    inputs = jnp.ones((1, 2, 1, 4))
+    position = jnp.array([[0, 1]])
+    outputs = layer(inputs, position=position)
+    self.assertEqual(outputs.shape, (1, 2, 1, 4))
+
+    # Snapshot verification
+    expected = jnp.array([[[[1.0, 1.0, 1.0, 1.0]], [[-0.300781, 0.996094, 1.38281, 1.00781]]]])
+    np.testing.assert_allclose(outputs, expected, atol=1e-5)
+
+
+if __name__ == "__main__":
+  unittest.main()

tests/unit/partial_rotary_embedding_test.py

@@ -163,6 +163,23 @@ def get_attn(pos):
         err_msg="PartialRotaryEmbedding attention should be shift-invariant.",
     )
 
+  def test_snapshot_verification(self):
+    """Verify output values against captured snapshot."""
+    layer = PartialRotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=4,
+        partial_rotary_factor=0.5,
+        rngs=self.nnx_rng,
+    )
+    inputs = jnp.ones((1, 2, 1, 4))
+    position = jnp.array([[0, 1]])
+    outputs = layer(inputs, position=position)
+
+    expected = jnp.array([[[[1.0, 1.0, 1.0, 1.0]], [[-0.30078125, 1.3828125, 1.0, 1.0]]]])
+    np.testing.assert_allclose(outputs, expected, atol=1e-5)
+
 
 class Gemma4PartialRotaryEmbeddingTest(unittest.TestCase):
   """Tests for the Gemma4PartialRotaryEmbedding layer."""
@@ -278,6 +295,23 @@ def get_attn(pos):
         err_msg="Gemma4PartialRotaryEmbedding attention should be shift-invariant.",
     )
 
+  def test_snapshot_verification(self):
+    """Verify output values against captured snapshot."""
+    layer = Gemma4PartialRotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=4,
+        partial_rotary_factor=0.5,
+        rngs=self.nnx_rng,
+    )
+    inputs = jnp.ones((1, 2, 1, 4))
+    position = jnp.array([[0, 1]])
+    outputs = layer(inputs, position=position)
+
+    expected = jnp.array([[[[1.0, 1.0, 1.0, 1.0]], [[-0.300781, 1.0, 1.38281, 1.0]]]])
+    np.testing.assert_allclose(outputs, expected, atol=1e-5)
+
 
 if __name__ == "__main__":
   unittest.main()