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

ai_edge_torch/generative/layers/unet/model_config.py

@@ -14,8 +14,7 @@
 # ==============================================================================
 
 # UNet configuration class.
-from dataclasses import dataclass
-from dataclasses import field
+import dataclasses
 import enum
 from typing import List, Optional
 
@@ -30,13 +29,13 @@ class SamplingType(enum.Enum):
   CONVOLUTION = enum.auto()
 
 
-@dataclass
+@dataclasses.dataclass
 class UpSamplingConfig:
   mode: SamplingType
   scale_factor: float
 
 
-@dataclass
+@dataclasses.dataclass
 class DownSamplingConfig:
   mode: SamplingType
   in_channels: int
@@ -46,7 +45,7 @@ class DownSamplingConfig:
   out_channels: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class ResidualBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -56,7 +55,7 @@ class ResidualBlock2DConfig:
   time_embedding_channels: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class AttentionBlock2DConfig:
   dim: int
   normalization_config: layers_cfg.NormalizationConfig
@@ -65,7 +64,7 @@ class AttentionBlock2DConfig:
   attention_batch_size: int = 1
 
 
-@dataclass
+@dataclasses.dataclass
 class CrossAttentionBlock2DConfig:
   query_dim: int
   cross_dim: int
@@ -75,7 +74,7 @@ class CrossAttentionBlock2DConfig:
   attention_batch_size: int = 1
 
 
-@dataclass
+@dataclasses.dataclass
 class FeedForwardBlock2DConfig:
   dim: int
   hidden_dim: int
@@ -84,7 +83,7 @@ class FeedForwardBlock2DConfig:
   use_bias: bool
 
 
-@dataclass
+@dataclasses.dataclass
 class TransformerBlock2DConfig:
   pre_conv_normalization_config: layers_cfg.NormalizationConfig
   attention_block_config: AttentionBlock2DConfig
@@ -92,7 +91,7 @@ class TransformerBlock2DConfig:
   feed_forward_block_config: FeedForwardBlock2DConfig
 
 
-@dataclass
+@dataclasses.dataclass
 class UpDecoderBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -113,7 +112,7 @@ class UpDecoderBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class SkipUpDecoderBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -136,7 +135,7 @@ class SkipUpDecoderBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class DownEncoderBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -157,7 +156,7 @@ class DownEncoderBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class MidBlock2DConfig:
   in_channels: int
   normalization_config: layers_cfg.NormalizationConfig
@@ -173,7 +172,7 @@ class MidBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class AutoEncoderConfig:
   """Configurations of encoder/decoder in the autoencoder model."""
 
@@ -210,7 +209,7 @@ class AutoEncoderConfig:
   mid_block_config: MidBlock2DConfig
 
 
-@dataclass
+@dataclasses.dataclass
 class DiffusionModelConfig:
   """Configurations of Diffusion model."""
 

ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+# ==============================================================================

ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-import json
-
 from ai_edge_quantizer import quantizer
 from ai_edge_torch.generative.quantize import quant_attrs
 from ai_edge_torch.generative.quantize import quant_recipe

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -25,7 +25,8 @@ class LayerQuantRecipe:
   """Quantization recipe for a single Edge Generative API layer (e.g. Attention).
 
   Generic layer-scoped quantization recipe that specifies how this layer should
-  be quantized by the Edge Generative API. This is applicable to layers implemented
+  be quantized by the Edge Generative API. This is applicable to layers
+  implemented
   in ai_edge_torch/generative/layers/. Combinations of attributes that are not
   supported during runtime will be detected when .verify() is called.
 
@@ -83,7 +84,8 @@ class LayerQuantRecipe:
 class GenerativeQuantRecipe:
   """Quantization recipe for a model composed of the Edge Generative API layers.
 
-  Some layers can be specified with different `LayerQuantRecipe` for each block by
+  Some layers can be specified with different `LayerQuantRecipe` for each block
+  by
   providing a dictionary keyed by the TransformerBlock index, e.g. attention
   and feedforward. For example,
 
@@ -102,11 +104,11 @@ class GenerativeQuantRecipe:
     default: The quantization recipe for global scope of the model.
     embedding: Recipe for the embedding table.
     attention: Recipe for the attention blocks. This could be specified with
-      different LayerQuantRecipe for each block by providing a dictionary
-      keyed by the TransformerBlock index.
+      different LayerQuantRecipe for each block by providing a dictionary keyed
+      by the TransformerBlock index.
     feedforward: Recipe for the feedforward layers. This could be specified with
-      different LayerQuantRecipe for each block by providing a dictionary
-      keyed by the TransformerBlock index.
+      different LayerQuantRecipe for each block by providing a dictionary keyed
+      by the TransformerBlock index.
   """
 
   default: Optional[LayerQuantRecipe] = None

ai_edge_torch/generative/quantize/quant_recipe_utils.py

@@ -16,7 +16,8 @@
 """Helper functions to construct custom quantization recipes.
 
 These are intended for more advanced users who want to configure their own
-quantization recipes. For pre-constructed recipes, use `quant_recipes.py` instead.
+quantization recipes. For pre-constructed recipes, use `quant_recipes.py`
+instead.
 
 Typical usage example:
 

ai_edge_torch/generative/test/test_experimental_ekv.py

@@ -14,24 +14,23 @@
 # ==============================================================================
 # A suite of tests to validate experimental external KV Cache layers and models.
 
-import unittest
-
 from ai_edge_torch.generative.examples.experimental.gemma import gemma
 from ai_edge_torch.generative.examples.experimental.phi import phi2
 from ai_edge_torch.generative.examples.experimental.tiny_llama import tiny_llama  # NOQA
 from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-import numpy as np
 import torch
 
+from tensorflow.python.platform import googletest
+
 
-class TestExternalKVLayers(unittest.TestCase):
+class TestExternalKVLayers(googletest.TestCase):
 
   def _get_test_config(
       self, num_layers, head_dim, num_query_groups, kv_cache_max_len
   ):
     attn_config = cfg.AttentionConfig(
-        num_heads=1, num_query_groups=num_query_groups
+        num_heads=1, head_dim=head_dim, num_query_groups=num_query_groups
     )
     config = cfg.ModelConfig(
         kv_cache_max_len=kv_cache_max_len,
@@ -117,7 +116,7 @@ class TestExternalKVLayers(unittest.TestCase):
     self.assertEqual(input_specs[1].arg.name, "kv_v_0")
 
 
-class TestExternalKVModels(unittest.TestCase):
+class TestExternalKVModels(googletest.TestCase):
 
   def test_can_build_gemma(self):
     gemma.define_and_run_2b(checkpoint_path=None, test_model=True)
@@ -130,4 +129,4 @@ class TestExternalKVModels(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/generative/test/{loader_test.py → test_loader.py}

@@ -16,15 +16,16 @@
 
 import os
 import tempfile
-import unittest
 
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
 from ai_edge_torch.generative.utilities import loader as loading_utils
 import safetensors.torch
 import torch
 
+from tensorflow.python.platform import googletest
 
-class TestLoader(unittest.TestCase):
+
+class TestLoader(googletest.TestCase):
   """Unit tests that check weight loader."""
 
   def test_load_safetensors(self):
@@ -78,4 +79,4 @@ class TestLoader(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/generative/test/test_model_conversion.py

@@ -14,9 +14,6 @@
 # ==============================================================================
 # Testing model conversion for a few gen-ai models.
 import copy
-import os
-import tempfile
-import unittest
 
 import ai_edge_torch
 from ai_edge_torch.generative.examples.gemma import gemma
@@ -27,22 +24,24 @@ from ai_edge_torch.testing import model_coverage
 import numpy as np
 import torch
 
+from tensorflow.python.platform import googletest
 
-class TestModelConversion(unittest.TestCase):
+
+class TestModelConversion(googletest.TestCase):
   """Unit tests that check for model conversion and correctness."""
 
   def test_toy_model_with_kv_cache(self):
     config = toy_model_with_kv_cache.get_model_config()
-    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config).eval()
     idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
         [10], dtype=torch.int64
     )
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       self.assertTrue(
           model_coverage.compare_tflite_torch(
               edge_model,
@@ -57,16 +56,16 @@ class TestModelConversion(unittest.TestCase):
   def test_toy_model_with_multi_batches(self):
     config = toy_model_with_kv_cache.get_model_config()
     config.batch_size = 2
-    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config).eval()
     idx, input_pos = torch.tensor([[1], [2]], dtype=torch.long), torch.tensor(
         [10], dtype=torch.int64
     )
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       self.assertTrue(
           model_coverage.compare_tflite_torch(
               edge_model,
@@ -81,16 +80,16 @@ class TestModelConversion(unittest.TestCase):
   def test_toy_model_with_kv_cache_with_hlfb(self):
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True
-    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config).eval()
     idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
         [10], dtype=torch.int64
     )
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       self.assertTrue(
           model_coverage.compare_tflite_torch(
               edge_model,
@@ -105,7 +104,7 @@ class TestModelConversion(unittest.TestCase):
   def test_tiny_llama(self):
     self.skipTest("b/338288901")
     config = tiny_llama.get_fake_model_config_for_test()
-    pytorch_model = tiny_llama.TinyLLamma(config)
+    pytorch_model = tiny_llama.TinyLLamma(config).eval()
 
     idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
     tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
@@ -114,9 +113,9 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       self.assertTrue(
           model_coverage.compare_tflite_torch(
               edge_model,
@@ -130,7 +129,7 @@ class TestModelConversion(unittest.TestCase):
 
   def test_tiny_llama_multisig(self):
     config = tiny_llama.get_fake_model_config_for_test()
-    pytorch_model = tiny_llama.TinyLLamma(config)
+    pytorch_model = tiny_llama.TinyLLamma(config).eval()
 
     # prefill
     seq_len = 10
@@ -151,9 +150,9 @@ class TestModelConversion(unittest.TestCase):
         .convert()
     )
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       copied_model = copy.deepcopy(pytorch_model)
 
       self.assertTrue(
@@ -188,9 +187,9 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       # TODO(talumbau, haoliang): debug numerical diff.
       self.assertTrue(
           model_coverage.compare_tflite_torch(
@@ -206,7 +205,7 @@ class TestModelConversion(unittest.TestCase):
   def test_phi2(self):
     self.skipTest("b/338288901")
     config = phi2.get_fake_model_config_for_test()
-    pytorch_model = phi2.Phi2(config)
+    pytorch_model = phi2.Phi2(config).eval()
 
     idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
     tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
@@ -215,9 +214,9 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
 
-    # TODO(b/338288901): re-enable test to check output tensors.
+    # TODO: b/338288901 - re-enable test to check output tensors.
     skip_output_check = True
-    if skip_output_check is False:
+    if not skip_output_check:
       self.assertTrue(
           model_coverage.compare_tflite_torch(
               edge_model,
@@ -231,4 +230,4 @@ class TestModelConversion(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/generative/test/test_quantize.py

@@ -13,9 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
-import unittest
-
 import ai_edge_torch
+from ai_edge_torch import config
 from ai_edge_torch.generative.examples.test_models import toy_model  # NOQA
 from ai_edge_torch.generative.quantize import quant_recipe
 from ai_edge_torch.generative.quantize import quant_recipe_utils
@@ -29,8 +28,10 @@ from ai_edge_torch.testing import model_coverage
 from parameterized import parameterized
 import torch
 
+from tensorflow.python.platform import googletest
+
 
-class TestVerifyRecipes(unittest.TestCase):
+class TestVerifyRecipes(googletest.TestCase):
   """Unit tests that check for model quantization recipes."""
 
   @parameterized.expand([
@@ -87,7 +88,7 @@ class TestVerifyRecipes(unittest.TestCase):
     ).verify()
 
 
-class TestQuantizeConvert(unittest.TestCase):
+class TestQuantizeConvert(googletest.TestCase):
   """Test conversion with quantization."""
 
   def _attention_int8_dynamic_recipe() -> quant_config.QuantConfig:
@@ -111,6 +112,10 @@ class TestQuantizeConvert(unittest.TestCase):
       (_attention_int8_dynamic_recipe()),
       (_feedforward_int8_dynamic_recipe()),
   ])
+  @googletest.skipIf(
+      not config.Config.use_torch_xla,
+      reason="Not working with odml_torch at the moment.",
+  )
   def test_quantize_convert_toy_sizes(self, quant_config):
     config = toy_model.get_model_config()
     pytorch_model = toy_model.ToySingleLayerModel(config)
@@ -157,4 +162,4 @@ class TestQuantizeConvert(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/generative/utilities/loader.py

@@ -92,9 +92,7 @@ def load_pytorch_statedict(full_path: str):
 
 
 class ModelLoader:
-  """A utility class for loading and converting model checkpoints to the
-  Edge Generative API layer format.
-  """
+  """Utlity for loading model checkpoints to the Edge Generative API layer."""
 
   @dataclass
   class TensorNames:
@@ -116,12 +114,13 @@ class ModelLoader:
     lm_head: str = None
 
   def __init__(self, file_name: str, names: TensorNames) -> None:
-    """ModelLoader constructor. Can be used to load multiple models of the same
-    type.
+    """ModelLoader constructor.
+
+    Can be used to load multiple models of the same type.
 
     Args:
-        file_name (str): Path to the checkpoint. Can be a directory or an
-          exact file.
+        file_name (str): Path to the checkpoint. Can be a directory or an exact
+          file.
         names (TensorNames): An instance of `TensorNames` to determine mappings.
     """
     self._file_name = file_name
@@ -140,7 +139,8 @@ class ModelLoader:
 
     Returns:
         missing_keys (List[str]): a list of str containing the missing keys.
-        unexpected_keys (List[str]): a list of str containing the unexpected keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected
+        keys.
 
     Raises:
         ValueError: If conversion results in unmapped tensors and strict mode is
@@ -208,7 +208,7 @@ class ModelLoader:
     if self._file_name.endswith(".bin") or self._file_name.endswith(".pt"):
       return load_pytorch_statedict
 
-    raise ValueError(f"File format not supported.")
+    raise ValueError("File format not supported.")
 
   def _map_feedforward(
       self,
@@ -346,9 +346,9 @@ class ModelLoader:
       q_per_kv = (
           config.attn_config.num_heads // config.attn_config.num_query_groups
       )
-      qs = torch.split(q, config.head_dim * q_per_kv)
-      ks = torch.split(k, config.head_dim)
-      vs = torch.split(v, config.head_dim)
+      qs = torch.split(q, config.attn_config.head_dim * q_per_kv)
+      ks = torch.split(k, config.attn_config.head_dim)
+      vs = torch.split(v, config.attn_config.head_dim)
       cycled = [t for group in zip(qs, ks, vs) for t in group]
       return torch.cat(cycled)
     else: