tpu-inference 0.12.0.dev20251213__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tests/kernels/gmm_test.py

@@ -0,0 +1,205 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+import jax
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
+from jax._src import test_util as jtu
+
+from tpu_inference.kernels.megablox.gmm import gmm
+
+jax.config.parse_flags_with_absl()
+
+
+def quantize_tensor(x: jax.Array,
+                    dtype: jnp.dtype,
+                    axis: int = -1,
+                    block_size: int = 256):
+    if jnp.issubdtype(dtype, jnp.integer):
+        dtype_info = jnp.iinfo(dtype)
+        max_val = int(dtype_info.max)
+        min_val = int(dtype_info.min)
+    else:
+        dtype_info = jnp.finfo(dtype)
+        max_val = float(dtype_info.max)
+        min_val = float(dtype_info.min)
+
+    orig_shape = x.shape
+    blocked_shape = orig_shape[:axis] + (-1,
+                                         block_size) + orig_shape[axis + 1:]
+    x_blocked = x.reshape(blocked_shape)
+
+    x_blocked_abs_max = jnp.max(jnp.abs(x_blocked),
+                                axis=axis + 1,
+                                keepdims=True)
+    scale = x_blocked_abs_max / max_val
+    x_blocked_q = jnp.clip(x_blocked / scale, min_val, max_val).astype(dtype)
+
+    x_q = x_blocked_q.reshape(orig_shape)
+    scale = scale.squeeze(axis=axis + 1).astype(jnp.float32)
+    return x_q, scale
+
+
+def reference_gmm(
+    lhs: jax.Array,
+    rhs: jax.Array,
+    group_sizes: jax.Array,
+    rhs_scale: jax.Array | None = None,
+    rhs_bias: jax.Array | None = None,
+    group_offset: jax.Array | None = None,
+):
+    num_groups, out_size, in_size = rhs.shape
+    assert lhs.shape[1] == in_size
+
+    if group_offset is None:
+        group_offset = jnp.array(0, dtype=jnp.int32)
+    start = group_sizes[:group_offset].sum()
+    group_sizes = group_sizes[group_offset:]
+    assert len(group_sizes) == num_groups
+
+    if rhs_scale is not None:
+        num_blocks = rhs_scale.shape[1]
+    else:
+        num_blocks = 1
+    block_size = in_size // num_blocks
+
+    gmm_out = [jnp.zeros((start, out_size), lhs.dtype)]
+    for group in range(num_groups):
+        end = start + group_sizes[group]
+
+        lhs_slice = lhs[start:end]
+        rhs_slice = rhs[group]
+
+        out = 0
+        for block in range(num_blocks):
+            block_start = block * block_size
+            block_end = block_start + block_size
+            lhs_block = lhs_slice[:, block_start:block_end].astype(jnp.float32)
+            rhs_block = rhs_slice[:, block_start:block_end].astype(jnp.float32)
+
+            acc = jnp.einsum("bd,hd->bh", lhs_block, rhs_block)
+            if rhs_scale is not None:
+                acc *= rhs_scale[group][block]
+            out += acc
+        if rhs_bias is not None:
+            out = out + rhs_bias[group]
+
+        gmm_out.append(out.astype(lhs.dtype))
+        start = end
+
+    return jnp.concat(gmm_out, axis=0)
+
+
+@jtu.with_config(jax_numpy_dtype_promotion="standard")
+class GmmTest(jtu.JaxTestCase):
+
+    @parameterized.product(
+        batch_size=[128],
+        in_size=[1024],
+        out_size=[1024],
+        num_groups=[16, 32],
+        has_bias=[True, False],
+    )
+    def test_gmm(self, batch_size, in_size, out_size, num_groups, has_bias):
+        key = jax.random.key(0)
+
+        lhs = jax.random.normal(key, (batch_size, in_size), dtype=jnp.bfloat16)
+        rhs = jax.random.normal(key, (num_groups, out_size, in_size),
+                                dtype=jnp.bfloat16)
+        rhs_bias = None
+        if has_bias:
+            rhs_bias = jax.random.normal(key, (num_groups, 1, out_size),
+                                         dtype=jnp.bfloat16)
+
+        group_sizes = jax.random.randint(key, (num_groups, ),
+                                         0,
+                                         batch_size,
+                                         dtype=jnp.int32)
+
+        expected = reference_gmm(lhs, rhs, group_sizes, rhs_bias=rhs_bias)
+        actual = gmm(
+            lhs,
+            rhs,
+            group_sizes,
+            rhs_bias=rhs_bias,
+            transpose_rhs=True,
+            preferred_element_type=jnp.bfloat16,
+        )
+
+        self.assertArraysAllClose(actual, expected)
+
+    @parameterized.product(
+        batch_size=[128],
+        in_size=[1024],
+        out_size=[1024],
+        num_groups=[16, 32],
+        has_bias=[True, False],
+        weight_dtype=[jnp.int8, jnp.float8_e5m2, jnp.float4_e2m1fn],
+        block_size=[256, 512],
+    )
+    def test_gmm_weight_quantized(
+        self,
+        batch_size,
+        in_size,
+        out_size,
+        num_groups,
+        has_bias,
+        weight_dtype,
+        block_size,
+    ):
+        if weight_dtype == jnp.float4_e2m1fn and not jtu.is_device_tpu_at_least(
+                version=7):
+            self.skipTest("Expect TPUv7+")
+        key = jax.random.key(0)
+
+        lhs = jax.random.normal(key, (batch_size, in_size), dtype=jnp.bfloat16)
+        rhs = jax.random.normal(key, (num_groups, out_size, in_size),
+                                dtype=jnp.bfloat16)
+        rhs_q, rhs_scale = quantize_tensor(rhs,
+                                           weight_dtype,
+                                           axis=2,
+                                           block_size=block_size)
+        rhs_scale = jnp.swapaxes(rhs_scale, 1, 2)
+        rhs_scale = jnp.expand_dims(rhs_scale, axis=2)
+
+        rhs_bias = None
+        if has_bias:
+            rhs_bias = jax.random.normal(key, (num_groups, 1, out_size),
+                                         dtype=jnp.bfloat16)
+
+        group_sizes = jax.random.randint(key, (num_groups, ),
+                                         0,
+                                         batch_size,
+                                         dtype=jnp.int32)
+
+        expected = reference_gmm(lhs,
+                                 rhs_q,
+                                 group_sizes,
+                                 rhs_scale=rhs_scale,
+                                 rhs_bias=rhs_bias)
+        actual = gmm(
+            lhs,
+            rhs_q,
+            group_sizes,
+            rhs_scale=rhs_scale,
+            rhs_bias=rhs_bias,
+            transpose_rhs=True,
+            preferred_element_type=jnp.bfloat16,
+        )
+
+        self.assertArraysAllClose(actual, expected, atol=3e-1, rtol=3e-1)
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tests/kernels/mla_v1_test.py

@@ -1,3 +1,17 @@
+# Copyright 2025 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
+#
+#     http://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.
+
 import jax
 import jax.numpy as jnp
 import numpy as np

tests/kernels/ragged_kv_cache_update_v2_test.py

@@ -1,3 +1,17 @@
+# Copyright 2025 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
+#
+#     http://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.
+
 import jax
 import jax.numpy as jnp
 import numpy as np

tests/kernels/ragged_paged_attention_kernel_v2_test.py

@@ -1,3 +1,17 @@
+# Copyright 2025 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
+#
+#     http://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.
+
 import random
 
 import jax

tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py

@@ -1,3 +1,17 @@
+# Copyright 2025 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
+#
+#     http://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.
+
 import jax
 import jax.numpy as jnp
 import numpy as np

tests/kernels/ragged_paged_attention_kernel_v3_test.py

@@ -1,3 +1,17 @@
+# Copyright 2025 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
+#
+#     http://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.
+
 import jax
 import jax.numpy as jnp
 import numpy as np

tests/layers/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+#     http://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/layers/common/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+#     http://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/layers/common/test_attention_interface.py

@@ -0,0 +1,156 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+from unittest.mock import MagicMock
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from jax.sharding import Mesh
+
+from tpu_inference.layers.common.attention_interface import attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.runner.kv_cache import get_kv_cache_shape_with_mesh
+
+# ---- Test Configuration & Constants ----
+
+# Total number of tokens across all sequences in the batch
+TOTAL_TOKENS = 10
+# Number of sequences in the batch
+NUM_SEQS = 2
+# Padded maximum number of sequences
+MAX_NUM_SEQS = 4
+# Number of attention heads (Query)
+NUM_HEADS = 8
+# Number of attention heads (Key/Value) - for Grouped-Query Attention
+NUM_KV_HEADS = 4
+# Total number of blocks in the KV cache
+NUM_BLOCKS = 32
+# Number of tokens per block
+BLOCK_SIZE = 16
+# Maximum number of blocks a single sequence can occupy
+MAX_BLOCKS_PER_SEQ = 8
+
+
+@pytest.fixture
+def mesh():
+    """Provides a mock 1D JAX mesh for testing."""
+    # Create a mesh with available devices, useful for running on CPU/GPU/TPU
+    # For this test, it will likely be a single CPU device.
+    devices = np.array(jax.local_devices()[:1])
+    if not devices.any():
+        # Add a mock device if no devices are present (e.g., in a CI environment)
+        devices = np.array([jax.devices("cpu")[0]])
+    return Mesh(devices.reshape((-1, 1, 1)), ("data", "attn_dp", "model"))
+
+
+# ---- Test for `attention` ----
+
+
+def _test_attention(monkeypatch, mesh, head_dim, use_sinks=False):
+    """
+    Tests the main `attention` function.
+
+    Verifies that:
+    1. It calls the `sharded_ragged_paged_attention` kernel with correct metadata.
+    2. The final outputs (kv_cache and attention output) have the correct shapes.
+    """
+    # 1. Arrange
+
+    # Create input tensors
+    q_dtype = jnp.float32
+    kv_dtype = jnp.float32
+    q = jnp.ones((TOTAL_TOKENS, NUM_HEADS, head_dim), dtype=q_dtype)
+    k = jnp.ones((TOTAL_TOKENS, NUM_KV_HEADS, head_dim), dtype=kv_dtype)
+    v = jnp.ones((TOTAL_TOKENS, NUM_KV_HEADS, head_dim), dtype=kv_dtype)
+    sinks = jnp.ones((NUM_HEADS, ), dtype=jnp.float32) if use_sinks else None
+
+    kv_cache_shape = get_kv_cache_shape_with_mesh(
+        mesh,
+        NUM_BLOCKS,
+        BLOCK_SIZE,
+        NUM_KV_HEADS,
+        head_dim,
+        kv_dtype,
+    )
+    kv_cache = jnp.zeros(kv_cache_shape, dtype=kv_dtype)
+
+    # Mock ragged_paged_attention to return a tensor of the correct shape
+    mock_paged_attn_kernel = MagicMock(return_value=(jnp.ones(
+        (TOTAL_TOKENS, NUM_HEADS, head_dim)), kv_cache), )
+
+    if head_dim == 64:
+        monkeypatch.setattr(
+            "tpu_inference.layers.common.attention_interface.ragged_paged_attention_hd64",
+            mock_paged_attn_kernel,
+        )
+    else:
+        monkeypatch.setattr(
+            "tpu_inference.layers.common.attention_interface.ragged_paged_attention",
+            mock_paged_attn_kernel,
+        )
+
+    # Create AttentionMetadata
+    attention_metadata = AttentionMetadata(
+        input_positions=jnp.arange(TOTAL_TOKENS, dtype=jnp.int32),
+        block_tables=jnp.zeros((MAX_NUM_SEQS * MAX_BLOCKS_PER_SEQ, ),
+                               dtype=jnp.int32),
+        seq_lens=jnp.array([5, 5, 0, 0], dtype=jnp.int32),
+        query_start_loc=jnp.array([0, 5, 10, 10, 10], dtype=jnp.int32),
+        request_distribution=jnp.array([0, 0, NUM_SEQS], dtype=jnp.int32),
+    )
+
+    # 2. Act
+    final_kv_cache, output = attention(
+        kv_cache=kv_cache,
+        q=q,
+        k=k,
+        v=v,
+        attention_metadata=attention_metadata,
+        mesh=mesh,
+        head_dim_original=head_dim,
+        sinks=sinks,
+    )
+
+    # 3. Assert
+    # Check that both mocked kernels were called
+    mock_paged_attn_kernel.assert_called_once()
+
+    # Check output shapes
+    assert final_kv_cache.shape == kv_cache.shape
+    assert output.shape == q.shape
+
+    # Check that the output is the one from our mock
+    assert jnp.all(output == 1.0)
+
+
+def test_attention(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 128)
+
+
+def test_attention_hd64(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 64)
+
+
+def test_attention_sink(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 64, True)
+
+
+def test_attention_sink_no_64_raises_error(monkeypatch, mesh):
+    with pytest.raises(
+            NotImplementedError,
+            match="Attention sink support is only available when head_dim==64"
+    ):
+        _test_attention(monkeypatch, mesh, 128, True)

tests/layers/common/test_quantization.py

@@ -0,0 +1,149 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+import functools
+
+import jax
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
+from jax._src import test_util as jtu
+
+from tpu_inference.layers.common.quantization import (
+    dequantize_tensor, dequantize_tensor_from_mxfp4_packed, quantize_kv,
+    quantize_tensor, quantize_tensor_to_mxfp4_packed)
+
+
+@jtu.with_config(jax_numpy_dtype_promotion="standard")
+class QuantizationTest(jtu.JaxTestCase):
+
+    @parameterized.product(axis=[-1, 0, (0, 1)])
+    def test_mxfp4_quantization(self, axis):
+        if not jtu.is_device_tpu_at_least(version=7):
+            self.skipTest("mxfp4 is only supported in TPUv7+")
+
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor_to_mxfp4_packed(original, axis)
+        dequantized = dequantize_tensor_from_mxfp4_packed(
+            tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.5, atol=0.5)
+
+    @parameterized.product(dtype=[jnp.float8_e4m3fn, jnp.int8],
+                           axis=[None, -1, 1, (0, 1)])
+    def test_quantization(self, dtype, axis):
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor(dtype, original, axis)
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.1, atol=0.1)
+
+    @parameterized.product(dtype=[jnp.float8_e4m3fn, jnp.int8],
+                           axis=[-1, 1],
+                           block_size=[32, 64])
+    def test_block_quantization(self, dtype, axis, block_size):
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor(dtype, original, axis, block_size)
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.1, atol=0.1)
+
+    @parameterized.product(dtype=[jnp.float8_e4m3fn, jnp.int8],
+                           axis=[(0, 1), (-1, 0)],
+                           block_size=[32, (64, 32)])
+    def test_multi_block_quantization(self, dtype, axis, block_size):
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor(dtype, original, axis, block_size)
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.1, atol=0.1)
+
+    def test_unaligned_block_quantization_raises_error(self):
+        key = jax.random.key(0)
+
+        shape = (128, 128)
+        tensor = jax.random.normal(key, shape, jnp.bfloat16)
+        block_size = 100
+        axis = 0
+
+        self.assertRaises(
+            ValueError,
+            functools.partial(quantize_tensor, jnp.int8, tensor, axis,
+                              block_size))
+
+    def test_block_quantization_padding(self):
+        key = jax.random.key(0)
+
+        shape = (128, 128)
+
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+        block_size = 100
+        axis = 0
+
+        tensor_q, scale = quantize_tensor(jnp.int8, original, axis, block_size,
+                                          True)
+
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        padded_size = ((shape[axis] + block_size) // block_size) * block_size
+        self.assertEqual(tensor_q.shape[axis], padded_size)
+        self.assertTrue((tensor_q[shape[0]:] == 0).all())
+        self.assertAllClose(dequantized[:shape[0]],
+                            original,
+                            rtol=0.1,
+                            atol=0.1)
+
+    @parameterized.product(kv_quant_dtype=[jnp.float8_e4m3fn, jnp.int8])
+    def test_quantize_kv(self, kv_quant_dtype):
+        """Tests the quantize_kv function with float8_e4m3fn dtype."""
+        key = jax.random.key(0)
+
+        shape = (128, 128)
+        k_original = jax.random.normal(key, shape, jnp.bfloat16)
+        v_original = jax.random.normal(key, shape, jnp.bfloat16)
+        k_scale = 0.1
+        v_scale = 0.2
+
+        k_quantized, v_quantized = quantize_kv(
+            kv_quant_dtype,
+            k_original,
+            v_original,
+            k_scale,
+            v_scale,
+        )
+
+        k_dequantized = k_quantized.astype(jnp.bfloat16) * k_scale
+        v_dequantized = v_quantized.astype(jnp.bfloat16) * v_scale
+
+        self.assertAllClose(k_dequantized, k_original, rtol=0.2, atol=0.2)
+        self.assertAllClose(v_dequantized, v_original, rtol=0.2, atol=0.2)
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tests/layers/jax/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+#     http://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/layers/jax/attention/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+#     http://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.