ai-edge-torch-nightly 0.2.0.dev20240806__py3-none-any.whl → 0.3.0.dev20240809__py3-none-any.whl

ai_edge_torch/{convert → _convert}/test/test_convert.py

@@ -12,24 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests for ai_edge_torch.convert."""
 
-
-from dataclasses import dataclass
+import dataclasses
 import os
-import tempfile
 from typing import Tuple
-import unittest
 
 import ai_edge_torch
-from ai_edge_torch.convert import conversion_utils as cutils
+from ai_edge_torch import config
+from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch.testing import model_coverage
 import numpy as np
 import tensorflow as tf
 import torch
 import torchvision
 
+from tensorflow.python.platform import googletest
+
 
-@dataclass
+@dataclasses.dataclass
 class TestContainer1:
   data_1: torch.Tensor
   data_2: Tuple[torch.Tensor, torch.Tensor]
@@ -40,10 +41,11 @@ torch.export.register_dataclass(
 )
 
 
-class TestConvert(unittest.TestCase):
+class TestConvert(googletest.TestCase):
   """Tests conversion of various modules."""
 
   def setUp(self):
+    super().setUp()
     torch.manual_seed(0)
 
   def test_convert_add(self):
@@ -66,8 +68,9 @@ class TestConvert(unittest.TestCase):
     )
 
   def test_convert_dot_add(self):
+    """Tests conversion of a matrix multiplication followed by an add."""
+
     class DotAdd(torch.nn.Module):
-      """Tests conversion of a matrix multiplication followed by an add."""
 
       def forward(self, a, b, c):
         return a @ b + c
@@ -97,20 +100,21 @@ class TestConvert(unittest.TestCase):
     """Tests conversion of a model with more than 10 arguments."""
 
     class AddChainWith11Args(torch.nn.Module):
+      """A model with 11 arguments."""
 
       def forward(
           self,
-          arg0: "f32[64]",
-          arg1: "f32[64]",
-          arg2: "f32[64]",
-          arg3: "f32[64]",
-          arg4: "f32[64]",
-          arg5: "f32[64]",
-          arg6: "f32[64]",
-          arg7: "f32[64]",
-          arg8: "f32[64]",
-          arg9: "f32[64]",
-          arg10: "f32[64]",
+          arg0: torch.Tensor,
+          arg1: torch.Tensor,
+          arg2: torch.Tensor,
+          arg3: torch.Tensor,
+          arg4: torch.Tensor,
+          arg5: torch.Tensor,
+          arg6: torch.Tensor,
+          arg7: torch.Tensor,
+          arg8: torch.Tensor,
+          arg9: torch.Tensor,
+          arg10: torch.Tensor,
       ):
         add0 = torch.add(arg0, arg1)
         add1 = torch.add(add0, arg2)
@@ -149,6 +153,7 @@ class TestConvert(unittest.TestCase):
     """Tests conversion of a model that returns multiple outputs."""
 
     class BasicAddModelWithMultipleOutputs(torch.nn.Module):
+      """A model that returns multiple outputs."""
 
       def forward(self, arg0, arg1):
         add0 = arg0 + arg1
@@ -172,6 +177,7 @@ class TestConvert(unittest.TestCase):
     """Tests conversion of a model that returns multiple outputs."""
 
     class BasicAddModelWithMultipleOutputs(torch.nn.Module):
+      """A model that returns multiple outputs."""
 
       def forward(self, arg0, arg1):
         add0 = arg0 + arg1
@@ -215,8 +221,8 @@ class TestConvert(unittest.TestCase):
     )
     self.assertTrue(result)
 
-  def test_apply_tfl_backdoor_flags(self):
-    """Tests if _apply_tfl_backdoor_flags correctly sets the values in a Converter object."""
+  def test_apply_tfl_converter_flags(self):
+    """Tests if _apply_tfl_converter_flags correctly sets the values in a Converter object."""
 
     class MockConverterInternalObject:
 
@@ -231,12 +237,12 @@ class TestConvert(unittest.TestCase):
 
     mock_converter = MockConverter()
     flags = {"key1": "new_value1", "key2": {"subkey2": "new_subvalue2"}}
-    cutils._apply_tfl_backdoor_flags(mock_converter, flags)
+    conversion_utils.apply_tfl_converter_flags(mock_converter, flags)
 
     self.assertTrue(flags["key1"], "new_value1")
     self.assertTrue(flags["key2"]["subkey2"], "new_subvalue2")
 
-  def test_convert_add_backdoor_flags(self):
+  def test_convert_add_converter_flags(self):
     """Tests conversion of an add module setting a tflite converter flag."""
 
     class Add(torch.nn.Module):
@@ -250,21 +256,23 @@ class TestConvert(unittest.TestCase):
     )
     torch_module = Add().eval()
 
-    with tempfile.TemporaryDirectory() as tmp_dir_path:
-      ir_dump_path = os.path.join(
-          tmp_dir_path, "test_convert_add_backdoor_flags_mlir_dump"
-      )
-      ai_edge_torch.convert(
-          torch_module,
-          args,
-          _ai_edge_converter_flags={"ir_dump_dir": ir_dump_path},
-      )
-      self.assertTrue(os.path.isdir(ir_dump_path))
+    tmp_dir_path = self.create_tempdir()
+    ir_dump_path = os.path.join(
+        tmp_dir_path, "test_convert_add_converter_flags_mlir_dump"
+    )
+    ai_edge_torch.convert(
+        torch_module,
+        args,
+        _ai_edge_converter_flags={"ir_dump_dir": ir_dump_path},
+    )
+    self.assertTrue(os.path.isdir(ir_dump_path))
 
+  @googletest.skipIf(
+      not config.Config.use_torch_xla,
+      reason="Shape polymorphism is not yet support with odml_torch.",
+  )
   def test_convert_model_with_dynamic_batch(self):
-    """
-    Test converting a simple model with dynamic batch size.
-    """
+    """Test converting a simple model with dynamic batch size."""
 
     class SampleModel(torch.nn.Module):
 
@@ -294,9 +302,7 @@ class TestConvert(unittest.TestCase):
       )
 
   def test_convert_model_with_kwargs(self):
-    """
-    Test converting a simple model with sample_kwargs.
-    """
+    """Test converting a simple model with sample_kwargs."""
 
     class SampleModel(torch.nn.Module):
 
@@ -315,9 +321,7 @@ class TestConvert(unittest.TestCase):
     )
 
   def test_convert_model_with_args_kwargs(self):
-    """
-    Test converting a simple model with both sample_args and sample_kwargs.
-    """
+    """Test converting a simple model with both sample_args and sample_kwargs."""
 
     class SampleModel(torch.nn.Module):
 
@@ -337,9 +341,7 @@ class TestConvert(unittest.TestCase):
     )
 
   def test_convert_model_with_args_nested_kwargs_1(self):
-    """
-    Test converting a simple model with both sample_args and nested sample_kwargs.
-    """
+    """Test converting a simple model with both sample_args and nested sample_kwargs."""
 
     class SampleModel(torch.nn.Module):
 
@@ -366,9 +368,7 @@ class TestConvert(unittest.TestCase):
     )
 
   def test_convert_model_with_args_nested_kwargs_2(self):
-    """
-    Test converting a simple model with both sample_args and nested sample_kwargs.
-    """
+    """Test converting a simple model with both sample_args and nested sample_kwargs."""
 
     class SampleModel(torch.nn.Module):
 
@@ -395,9 +395,7 @@ class TestConvert(unittest.TestCase):
     )
 
   def test_convert_model_with_args_nested_kwargs_3(self):
-    """
-    Test converting a simple model with both sample_args and nested sample_kwargs.
-    """
+    """Test converting a simple model with both sample_args and nested sample_kwargs."""
 
     class SampleModel(torch.nn.Module):
 
@@ -437,4 +435,4 @@ class TestConvert(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/{convert → _convert}/test/test_convert_composites.py

@@ -12,18 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests conversion modules that are meant to be wrapped as composites."""
 
-
-from typing import Callable
-import unittest
+from collections.abc import Callable
 
 import ai_edge_torch
 from ai_edge_torch.testing import model_coverage
 import parameterized
 import torch
 
+from tensorflow.python.platform import googletest
+
+
+def _func_to_torch_module(func: Callable[..., torch.Tensor]):
+  """Wraps a function into a torch module."""
 
-def _func_to_torch_module(func: Callable):
   class TestModule(torch.nn.Module):
 
     def __init__(self, func):
@@ -36,7 +39,7 @@ def _func_to_torch_module(func: Callable):
   return TestModule(func).eval()
 
 
-class TestConvertComposites(unittest.TestCase):
+class TestConvertComposites(googletest.TestCase):
   """Tests conversion modules that are meant to be wrapped as composites."""
 
   def test_convert_hardswish(self):
@@ -51,7 +54,8 @@ class TestConvertComposites(unittest.TestCase):
     )
 
   @parameterized.parameterized.expand([
-      # input_size, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override
+      # (input_size, kernel_size, stride, padding, ceil_mode,
+      # count_include_pad, divisor_override)
       # no padding, stride = 1
       ([1, 3, 6, 6], [3, 3], [1, 1], [0, 0], False, True, None),
       # add stride
@@ -64,6 +68,8 @@ class TestConvertComposites(unittest.TestCase):
       ([1, 3, 6, 6], [3, 3], [1, 1], [0, 1], False, True, None),
       # add both stride and padding
       ([1, 3, 6, 6], [3, 3], [2, 2], [1, 1], False, True, None),
+      # padding set to one number
+      ([1, 3, 6, 6], [3, 3], [1, 1], 1, False, True, None),
       # count_include_pad = False
       ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, False, None),
       # ceil_mode = True
@@ -72,8 +78,6 @@ class TestConvertComposites(unittest.TestCase):
       ([1, 3, 6, 6], [3, 3], [3, 3], [1, 1], True, True, None),
       # set divisor_override
       ([1, 3, 6, 6], [3, 3], [1, 1], 0, False, True, 6),
-      # padding set to one number
-      ([1, 3, 6, 6], [3, 3], [1, 1], 1, False, True, None),
   ])
   def test_convert_avg_pool2d(self, input_size, *args):
     """Tests conversion of a module containing an avg_pool2d aten."""
@@ -111,7 +115,7 @@ class TestConvertComposites(unittest.TestCase):
   def test_convert_upsample_bilinear_functional(self, input_size, kwargs):
     """Tests conversion of a torch.nn.functional.upsample module."""
     torch_module = _func_to_torch_module(
-        lambda input_tensor: torch.nn.functional.upsample(
+        lambda input_tensor: torch.nn.functional.upsample(  # pylint: disable=unnecessary-lambda
             input_tensor, **kwargs
         )
     )
@@ -146,7 +150,7 @@ class TestConvertComposites(unittest.TestCase):
   def test_convert_upsample_bilinear(self, input_size, kwargs):
     """Tests conversion of a torch.nn.Upsample module."""
     torch_module = _func_to_torch_module(
-        lambda input_tensor: torch.nn.Upsample(**kwargs)(input_tensor)
+        lambda input_tensor: torch.nn.Upsample(**kwargs)(input_tensor)  # pylint: disable=unnecessary-lambda
     )
     tracing_args = (torch.randn(*input_size),)
     edge_model = ai_edge_torch.convert(torch_module, tracing_args)
@@ -179,7 +183,7 @@ class TestConvertComposites(unittest.TestCase):
   def test_convert_interpolate_bilinear_functional(self, input_size, kwargs):
     """Tests conversion of a torch.nn.functional.interpolate module."""
     torch_module = _func_to_torch_module(
-        lambda input_tensor: torch.nn.functional.interpolate(
+        lambda input_tensor: torch.nn.functional.interpolate(  # pylint: disable=unnecessary-lambda
             input_tensor, **kwargs
         )
     )
@@ -227,4 +231,4 @@ class TestConvertComposites(unittest.TestCase):
 
 
 if __name__ == '__main__':
-  unittest.main()
+  googletest.main()

ai_edge_torch/{convert → _convert}/test/test_convert_multisig.py

@@ -12,19 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import unittest
+"""Tests for multi-signature conversion."""
 
 import ai_edge_torch
 from ai_edge_torch.testing import model_coverage
 import torch
 import torchvision
 
+from tensorflow.python.platform import googletest
+
 
-class TestConvertMultiSignature(unittest.TestCase):
+class TestConvertMultiSignature(googletest.TestCase):
   """Tests conversion of various modules through multi-signature conversion."""
 
   def setUp(self):
+    super().setUp()
     torch.manual_seed(0)
 
   def test_convert_mobilenet_v2_with_default(self):
@@ -144,4 +146,4 @@ class TestConvertMultiSignature(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/{convert → _convert}/test/test_to_channel_last_io.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-import unittest
+"""Tests for to_channel_last_io API and module wrapper."""
 
 import ai_edge_torch
 import torch
 
+from tensorflow.python.platform import googletest
+
 
 class Identity(torch.nn.Module):
 
@@ -25,7 +26,7 @@ class Identity(torch.nn.Module):
     return x
 
 
-class TestToChannelLastIO(unittest.TestCase):
+class TestToChannelLastIO(googletest.TestCase):
   """Tests to_channel_last_io API and module wrapper."""
 
   def test_no_transformations(self):
@@ -92,4 +93,4 @@ class TestToChannelLastIO(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/{convert → _convert}/to_channel_last_io.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Transforms the input and output of a module to channel last layout."""
 
 from typing import Optional
 
@@ -82,8 +83,10 @@ def to_channel_last_io(
       (N, C, ...) to channel last (N, ..., C).
     outputs (list[int]): Transform outputs with indices in the list from channel
       first (N, C, ...) to channel last (N, ..., C).
+
   Returns:
-    The wrapped nn.Module with additional layout transposes after inputs and/or before
+    The wrapped nn.Module with additional layout transposes after inputs and/or
+    before
     outputs.
   """
   return ChannelLastIOWrapper(module, args=args, outputs=outputs)

ai_edge_torch/config.py

@@ -0,0 +1,24 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Provides a configuration for the AI Edge Torch library."""
+
+import dataclasses
+import os
+
+
+@dataclasses.dataclass
+class Config:
+  use_torch_xla: bool = os.environ.get("USE_TORCH_XLA", "True") == "True"

ai_edge_torch/conftest.py

@@ -0,0 +1,20 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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 absl import flags
+
+
+def pytest_configure(config):
+  flags.FLAGS.mark_as_parsed()

ai_edge_torch/debug/culprit.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Culprit finder for AI Edge Torch conversion."""
 
 import contextlib
 import copy
@@ -20,14 +21,13 @@ import functools
 import io
 import operator
 import os
-import sys
 from typing import Any, Callable, Generator, List, Optional, Tuple, Union
 
 import ai_edge_torch
 from ai_edge_torch.debug import utils
-from functorch.compile import minifier as fx_minifier
 import torch
 from torch._functorch import aot_autograd
+from torch._functorch.fx_minifier import minifier as fx_minifier
 import torch.utils._pytree as pytree
 
 _torch_float_dtypes = {
@@ -116,7 +116,7 @@ class Culprit(SearchResult):
       print_output: bool - If true, prints the code to stdout. Otherwise returns
         the code in a str.
     """
-    # TODO (b/321263453): Support Python code gen with sample arg tensor values.
+    # TODO: b/321263453 - Support Python code gen with sample arg tensor values.
     random_inputs = True
 
     graph_module_code = self.graph_module.print_readable(
@@ -152,6 +152,7 @@ class Culprit(SearchResult):
 
   def print_code(self, print_output=True):
     """Print the Python code for culprit graph module, sample args, and AI
+
     Edge Torch conversion that will fail with the error.
 
     Args:
@@ -188,8 +189,8 @@ class Culprit(SearchResult):
 
 
 def _normalize_getitem_nodes(fx_gm: torch.fx.GraphModule):
-  """
-  This function turns all operator getitem nodes in ExportedProgram FX graph to
+  """This function turns all operator getitem nodes in ExportedProgram FX graph to
+
   new nodes composed of "computation + getitem". The normalization duplicates
   some computations in the graph but would make the graph more friendly for
   partitioning in FX minifier.
@@ -367,19 +368,18 @@ def _search_model(
     max_granularity: Optional[int] = None,
     enable_fx_minifier_logging: bool = False,
 ) -> Generator[SearchResult, None, None]:
-  """Finds subgraphs in the torch model that satisfy a certain predicate function provided by the users.
+  """Finds subgraphs in the torch model that satify a certain predicate function provided by the users.
 
   Args:
-    predicate_f: a predicate function the users specify.
-      It takes a FX (sub)graph and the inputs to this graph,
-      return True if the graph satisfies the predicate,
-      return False otherwise.
+    predicate_f: a predicate function the users specify. It takes a FX
+      (sub)graph and the inputs to this graph, return True if the graph
+      satisfies the predicate, return False otherwise.
     model: model in which to search subgraph.
-    export_args: A set of args to trace the model with,
-      i.e. model(*args) must run.
-    max_granularity - FX minifier arg. The maximum granularity (number of nodes)
-      in the returned ATen FX subgraph of the culprit.
-    enable_fx_minifier_logging: If true, allows the underlying FX minifier to log the progress.
+    export_args: A set of args to trace the model with, i.e. model(*args) must
+      run. max_granularity - FX minifier arg. The maximum granularity (number of
+      nodes) in the returned ATen FX subgraph of the culprit.
+    enable_fx_minifier_logging: If true, allows the underlying FX minifier to
+      log the progress.
   """
 
   if isinstance(model, torch.nn.Module):
@@ -469,13 +469,13 @@ def find_culprits(
 
   Args:
     torch_model: model to export and save
-    args: A set of args to trace the model with, i.e.
-      torch_model(*args) must run
-    max_granularity - FX minifier arg. The maximum granularity (number of nodes)
-      in the returned ATen FX subgraph of the culprit.
-    runtime_errors: If true, find culprits for Python runtime errors
-      with converted model.
-    enable_fx_minifier_logging: If true, allows the underlying FX minifier to log the progress.
+    args: A set of args to trace the model with, i.e. torch_model(*args) must
+      run max_granularity - FX minifier arg. The maximum granularity (number of
+      nodes) in the returned ATen FX subgraph of the culprit.
+    runtime_errors: If true, find culprits for Python runtime errors with
+      converted model.
+    enable_fx_minifier_logging: If true, allows the underlying FX minifier to
+      log the progress.
   """
 
   fx_minifier_checker = functools.partial(

ai_edge_torch/debug/test/test_culprit.py

@@ -17,11 +17,12 @@
 import ast
 import io
 import sys
-import unittest
 
 from ai_edge_torch.debug import find_culprits
 import torch
 
+from tensorflow.python.platform import googletest
+
 _test_culprit_lib = torch.library.Library("test_culprit", "DEF")
 
 _test_culprit_lib.define("non_lowerable_op(Tensor x) -> Tensor")
@@ -49,7 +50,7 @@ class BadModel(torch.nn.Module):
     return x
 
 
-class TestCulprit(unittest.TestCase):
+class TestCulprit(googletest.TestCase):
 
   def test_find_culprits(self):
     model = BadModel().eval()
@@ -131,4 +132,4 @@ class TestCulprit(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/debug/test/test_search_model.py

@@ -12,15 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-
-import unittest
+"""Tests for search_model."""
 
 from ai_edge_torch.debug import _search_model
 import torch
 
+from tensorflow.python.platform import googletest
+
 
-class TestSearchModel(unittest.TestCase):
+class TestSearchModel(googletest.TestCase):
 
   def test_search_model_with_ops(self):
     class MultipleOpsModel(torch.nn.Module):
@@ -48,4 +48,4 @@ class TestSearchModel(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/debug/utils.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Utils for debugging."""
+
 import contextlib
 import sys
 
 import torch
-from torch.export.graph_signature import InputKind
-import torch.fx._pytree as fx_pytree
 from torch.utils import _pytree as pytree
 
 
@@ -33,6 +33,15 @@ def exported_program_to_fx_graph_module_and_inputs(
 
 @contextlib.contextmanager
 def redirect_stdio(stdout, stderr):
+  """Redirects stdout and stderr to the given file objects.
+
+  Args:
+    stdout: A file object to redirect stdout to.
+    stderr: A file object to redirect stderr to.
+
+  Yields:
+    The file objects that stdout and stderr were redirected to.
+  """
   old_stdout = sys.stdout
   old_stderr = sys.stderr
 

ai_edge_torch/generative/examples/experimental/gemma/convert_to_tflite.py

@@ -34,8 +34,8 @@ def convert_gemma_to_tflite(
     quantize: bool = True,
 ):
   """An example method for converting a Gemma 2B model to multi-signature
-  tflite model.
 
+  tflite model.
   Args:
       checkpoint_path (str): The filepath to the model checkpoint, or directory
         holding the checkpoint.
@@ -43,8 +43,8 @@ def convert_gemma_to_tflite(
         Defaults to 512.
       kv_cache_max_len (int, optional): The maximum size of KV cache buffer,
         including both prefill and decode. Defaults to 1024.
-      quantize (bool, optional): Whether the model should be quanized.
-        Defaults to True.
+      quantize (bool, optional): Whether the model should be quanized. Defaults
+        to True.
   """
   pytorch_model = gemma.build_2b_model(
       checkpoint_path, kv_cache_max_len=kv_cache_max_len

ai_edge_torch/generative/examples/experimental/gemma/gemma.py

@@ -73,7 +73,9 @@ class Gemma(nn.Module):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -125,6 +127,7 @@ class Gemma(nn.Module):
 def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=8,
+      head_dim=256,
       num_query_groups=1,
       rotary_percentage=1.0,
   )