ai-edge-torch-nightly 0.2.0.dev20240801__py3-none-any.whl → 0.2.0.dev20240803__py3-none-any.whl

ai_edge_torch/generative/layers/attention.py

@@ -16,16 +16,15 @@
 
 from typing import Optional, Tuple
 
-import torch
-from torch import nn
-import torch.nn.functional as F
-
 import ai_edge_torch.generative.layers.builder as builder
 from ai_edge_torch.generative.layers.kv_cache import KVCache
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention  # NOQA
 from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention_with_hlfb  # NOQA
+import torch
+from torch import nn
+import torch.nn.functional as F
 
 
 def _embed_rope(
@@ -140,7 +139,9 @@ class CausalSelfAttention(nn.Module):
     shape = (config.num_heads + 2 * config.num_query_groups) * self.head_dim
     # Key, query, value projections for all heads.
     self.qkv_projection = nn.Linear(dim, shape, bias=config.qkv_use_bias)
-    self.output_projection = nn.Linear(dim, dim, bias=config.output_proj_use_bias)
+    self.output_projection = nn.Linear(
+        dim, dim, bias=config.output_proj_use_bias
+    )
     self.config = config
     self.kv_cache = None
     self.batch_size = batch_size
@@ -181,9 +182,10 @@ class CausalSelfAttention(nn.Module):
     """
     # Batch size, sequence length, embedding dimensionality.
     B, T, E = x.size()
-    assert (
-        B == self.batch_size
-    ), "batch size of input tensor must match with the batch size specified in the model configuration."
+    assert B == self.batch_size, (
+        "batch size of input tensor must match with the batch size specified in"
+        " the model configuration."
+    )
 
     qkv = self.qkv_projection(x)
 
@@ -279,9 +281,15 @@ class CrossAttention(nn.Module):
     self.config = config
     self.head_dim = query_dim // config.num_heads
     self.n_heads = config.num_heads
-    self.q_projection = nn.Linear(query_dim, query_dim, bias=config.qkv_use_bias)
-    self.k_projection = nn.Linear(cross_dim, query_dim, bias=config.qkv_use_bias)
-    self.v_projection = nn.Linear(cross_dim, query_dim, bias=config.qkv_use_bias)
+    self.q_projection = nn.Linear(
+        query_dim, query_dim, bias=config.qkv_use_bias
+    )
+    self.k_projection = nn.Linear(
+        cross_dim, query_dim, bias=config.qkv_use_bias
+    )
+    self.v_projection = nn.Linear(
+        cross_dim, query_dim, bias=config.qkv_use_bias
+    )
     self.output_projection = nn.Linear(
         query_dim, query_dim, bias=config.output_proj_use_bias
     )

ai_edge_torch/generative/layers/builder.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 # Builder class for individual components.
-import torch
-from torch import nn
-import torch.nn.functional as F
-
 import ai_edge_torch.generative.layers.feed_forward as feed_forward
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.normalization as normalization
+import torch
+from torch import nn
+import torch.nn.functional as F
 
 
 class GeGLU(nn.Module):

ai_edge_torch/generative/layers/kv_cache.py

@@ -14,16 +14,17 @@
 # ==============================================================================
 # `nn.Module` which implements a KV cache.
 
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 from torch import nn
 import torch_xla
 
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
-
 
 class KVCache(nn.Module):
 
-  def __init__(self, batch_size, kv_cache_max, n_heads, head_dim, enable_hlfb=False):
+  def __init__(
+      self, batch_size, kv_cache_max, n_heads, head_dim, enable_hlfb=False
+  ):
     """Initializes the KVCache layer.
 
     Args:

ai_edge_torch/generative/layers/model_config.py

@@ -124,9 +124,13 @@ class ModelConfig:
       default_factory=NormalizationConfig
   )
   # The normalization applied to feed forward's input.
-  pre_ff_norm_config: NormalizationConfig = field(default_factory=NormalizationConfig)
+  pre_ff_norm_config: NormalizationConfig = field(
+      default_factory=NormalizationConfig
+  )
   # The normalization applied before LM head.
-  final_norm_config: NormalizationConfig = field(default_factory=NormalizationConfig)
+  final_norm_config: NormalizationConfig = field(
+      default_factory=NormalizationConfig
+  )
 
   # If set to True, only pre_attention_norm is applied to the input and the
   # decode's output is computed as `output = input + attn_out + ff_out` where

ai_edge_torch/generative/layers/rotary_position_embedding.py

@@ -16,7 +16,9 @@
 import torch
 
 
-def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+def apply_rope(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
   """Computes rotary positional embedding.
 
   Args:

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -17,11 +17,10 @@
 import math
 from typing import Optional
 
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
 
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
-
 
 def scaled_dot_product_attention(
     q: torch.Tensor,

ai_edge_torch/generative/layers/unet/blocks_2d.py

@@ -15,15 +15,14 @@
 
 from typing import List, Optional, Tuple
 
-import torch
-from torch import nn
-
 from ai_edge_torch.generative.layers.attention import CrossAttention
 from ai_edge_torch.generative.layers.attention import SelfAttention
 import ai_edge_torch.generative.layers.builder as layers_builder
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 import ai_edge_torch.generative.layers.unet.builder as unet_builder
 import ai_edge_torch.generative.layers.unet.model_config as unet_cfg
+import torch
+from torch import nn
 
 
 class ResidualBlock2D(nn.Module):
@@ -41,7 +40,11 @@ class ResidualBlock2D(nn.Module):
         config.in_channels, config.normalization_config
     )
     self.conv_1 = nn.Conv2d(
-        config.in_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+        config.in_channels,
+        config.out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=1,
     )
     if config.time_embedding_channels is not None:
       self.time_emb_proj = nn.Linear(
@@ -53,14 +56,22 @@ class ResidualBlock2D(nn.Module):
         config.out_channels, config.normalization_config
     )
     self.conv_2 = nn.Conv2d(
-        config.out_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+        config.out_channels,
+        config.out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=1,
     )
     self.act_fn = layers_builder.get_activation(config.activation_config)
     if config.in_channels == config.out_channels:
       self.residual_layer = nn.Identity()
     else:
       self.residual_layer = nn.Conv2d(
-          config.in_channels, config.out_channels, kernel_size=1, stride=1, padding=0
+          config.in_channels,
+          config.out_channels,
+          kernel_size=1,
+          stride=1,
+          padding=0,
       )
 
   def forward(
@@ -105,7 +116,9 @@ class AttentionBlock2D(nn.Module):
     """
     super().__init__()
     self.config = config
-    self.norm = layers_builder.build_norm(config.dim, config.normalization_config)
+    self.norm = layers_builder.build_norm(
+        config.dim, config.normalization_config
+    )
     self.attention = SelfAttention(
         config.attention_batch_size,
         config.dim,
@@ -125,7 +138,10 @@ class AttentionBlock2D(nn.Module):
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
+    if (
+        self.config.normalization_config.type
+        == layers_cfg.NormalizationType.GROUP_NORM
+    ):
       x = self.norm(input_tensor)
       x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
@@ -156,7 +172,9 @@ class CrossAttentionBlock2D(nn.Module):
     """
     super().__init__()
     self.config = config
-    self.norm = layers_builder.build_norm(config.query_dim, config.normalization_config)
+    self.norm = layers_builder.build_norm(
+        config.query_dim, config.normalization_config
+    )
     self.attention = CrossAttention(
         config.attention_batch_size,
         config.query_dim,
@@ -180,7 +198,10 @@ class CrossAttentionBlock2D(nn.Module):
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
+    if (
+        self.config.normalization_config.type
+        == layers_cfg.NormalizationType.GROUP_NORM
+    ):
       x = self.norm(input_tensor)
       x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
@@ -209,7 +230,9 @@ class FeedForwardBlock2D(nn.Module):
     super().__init__()
     self.config = config
     self.act = layers_builder.get_activation(config.activation_config)
-    self.norm = layers_builder.build_norm(config.dim, config.normalization_config)
+    self.norm = layers_builder.build_norm(
+        config.dim, config.normalization_config
+    )
     if config.activation_config.type == layers_cfg.ActivationType.GE_GLU:
       self.w1 = nn.Identity()
       self.w2 = nn.Linear(config.hidden_dim, config.dim)
@@ -220,7 +243,10 @@ class FeedForwardBlock2D(nn.Module):
   def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
+    if (
+        self.config.normalization_config.type
+        == layers_cfg.NormalizationType.GROUP_NORM
+    ):
       x = self.norm(input_tensor)
       x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
@@ -287,7 +313,9 @@ class TransformerBlock2D(nn.Module):
         padding=0,
     )
     self.self_attention = AttentionBlock2D(config.attention_block_config)
-    self.cross_attention = CrossAttentionBlock2D(config.cross_attention_block_config)
+    self.cross_attention = CrossAttentionBlock2D(
+        config.cross_attention_block_config
+    )
     self.feed_forward = FeedForwardBlock2D(config.feed_forward_block_config)
     self.conv_out = nn.Conv2d(
         config.attention_block_config.dim,
@@ -371,7 +399,9 @@ class DownEncoderBlock2D(nn.Module):
       if config.transformer_block_config:
         transformers.append(TransformerBlock2D(config.transformer_block_config))
     self.resnets = nn.ModuleList(resnets)
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
     if config.add_downsample:
       self.downsampler = unet_builder.build_downsampling(config.sampling_config)
     else:
@@ -467,12 +497,18 @@ class UpDecoderBlock2D(nn.Module):
       if config.transformer_block_config:
         transformers.append(TransformerBlock2D(config.transformer_block_config))
     self.resnets = nn.ModuleList(resnets)
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
     if config.add_upsample:
       self.upsampler = unet_builder.build_upsampling(config.sampling_config)
       if config.upsample_conv:
         self.upsample_conv = nn.Conv2d(
-            config.out_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+            config.out_channels,
+            config.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
         )
     else:
       self.upsampler = None
@@ -548,9 +584,13 @@ class SkipUpDecoderBlock2D(nn.Module):
     transformers = []
     for i in range(config.num_layers):
       res_skip_channels = (
-          config.in_channels if (i == config.num_layers - 1) else config.out_channels
+          config.in_channels
+          if (i == config.num_layers - 1)
+          else config.out_channels
+      )
+      resnet_in_channels = (
+          config.prev_out_channels if i == 0 else config.out_channels
       )
-      resnet_in_channels = config.prev_out_channels if i == 0 else config.out_channels
       resnets.append(
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
@@ -565,12 +605,18 @@ class SkipUpDecoderBlock2D(nn.Module):
       if config.transformer_block_config:
         transformers.append(TransformerBlock2D(config.transformer_block_config))
     self.resnets = nn.ModuleList(resnets)
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
     if config.add_upsample:
       self.upsampler = unet_builder.build_upsampling(config.sampling_config)
       if config.upsample_conv:
         self.upsample_conv = nn.Conv2d(
-            config.out_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+            config.out_channels,
+            config.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
         )
     else:
       self.upsampler = None
@@ -678,7 +724,9 @@ class MidBlock2D(nn.Module):
       )
     self.resnets = nn.ModuleList(resnets)
     self.attentions = nn.ModuleList(attentions) if len(attentions) > 0 else None
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
 
   def forward(
       self,

ai_edge_torch/generative/layers/unet/builder.py

@@ -14,9 +14,8 @@
 # ==============================================================================
 # Builder utils for individual components.
 
-from torch import nn
-
 import ai_edge_torch.generative.layers.unet.model_config as unet_config
+from torch import nn
 
 
 def build_upsampling(config: unet_config.UpSamplingConfig):
@@ -30,10 +29,14 @@ def build_upsampling(config: unet_config.UpSamplingConfig):
 
 def build_downsampling(config: unet_config.DownSamplingConfig):
   if config.mode == unet_config.SamplingType.AVERAGE:
-    return nn.AvgPool2d(config.kernel_size, config.stride, padding=config.padding)
+    return nn.AvgPool2d(
+        config.kernel_size, config.stride, padding=config.padding
+    )
   elif config.mode == unet_config.SamplingType.CONVOLUTION:
     out_channels = (
-        config.in_channels if config.out_channels is None else config.out_channels
+        config.in_channels
+        if config.out_channels is None
+        else config.out_channels
     )
     padding = (0, 1, 0, 1) if config.padding == 0 else config.padding
     return nn.Conv2d(

ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py

@@ -16,7 +16,6 @@
 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
 
@@ -44,7 +43,9 @@ def _get_nbits_from_dtype(dtype: quant_attrs.Dtype) -> int:
   raise ValueError('Unimplemented number of bits')
 
 
-def _get_dtype_from_dtype(dtype: quant_attrs.Dtype) -> quantizer.qtyping.TensorDataType:
+def _get_dtype_from_dtype(
+    dtype: quant_attrs.Dtype,
+) -> quantizer.qtyping.TensorDataType:
   if dtype == quant_attrs.Dtype.FP32 or dtype == quant_attrs.Dtype.FP16:
     return quantizer.qtyping.TensorDataType.FLOAT
   else:
@@ -59,7 +60,9 @@ def _get_execution_mode_from_mode(mode: quant_attrs.Mode) -> _OpExecutionMode:
   raise ValueError('Unimplemented execution mode')
 
 
-def _get_channelwise_from_granularity(granularity: quant_attrs.Granularity) -> bool:
+def _get_channelwise_from_granularity(
+    granularity: quant_attrs.Granularity,
+) -> bool:
   if granularity == quant_attrs.Granularity.CHANNELWISE:
     return True
   elif granularity == quant_attrs.Granularity.NONE:
@@ -87,7 +90,9 @@ def _set_quant_config(
           weight_tensor_config=_TensorQuantConfig(
               num_bits=_get_nbits_from_dtype(layer_recipe.weight_dtype),
               symmetric=True,
-              channel_wise=_get_channelwise_from_granularity(layer_recipe.granularity),
+              channel_wise=_get_channelwise_from_granularity(
+                  layer_recipe.granularity
+              ),
               dtype=_get_dtype_from_dtype(layer_recipe.weight_dtype),
           ),
           execution_mode=_get_execution_mode_from_mode(layer_recipe.mode),

ai_edge_torch/generative/quantize/example.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-import numpy as np
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.gemma import gemma
 from ai_edge_torch.generative.quantize import quant_recipes
+import numpy as np
+import torch
 
 
 def main():

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -74,7 +74,8 @@ class LayerQuantRecipe:
 
     if not is_valid:
       raise ValueError(
-          'Unsupported LayerQuantRecipe configuration. See get_supported_recipe_matrix()'
+          'Unsupported LayerQuantRecipe configuration. See'
+          ' get_supported_recipe_matrix()'
       )
 
 

ai_edge_torch/generative/test/loader_test.py

@@ -18,11 +18,10 @@ import os
 import tempfile
 import unittest
 
-import safetensors.torch
-import torch
-
 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
 
 
 class TestLoader(unittest.TestCase):
@@ -59,7 +58,9 @@ class TestLoader(unittest.TestCase):
           "model.layers.0.mlp.down_proj.weight": torch.randn((2048, 5632)),
           "model.layers.0.mlp.gate_proj.weight": torch.randn((5632, 2048)),
           "model.layers.0.mlp.up_proj.weight": torch.randn((5632, 2048)),
-          "model.layers.0.post_attention_layernorm.weight": torch.randn((2048,)),
+          "model.layers.0.post_attention_layernorm.weight": torch.randn((
+              2048,
+          )),
           "model.layers.0.self_attn.k_proj.weight": torch.randn((256, 2048)),
           "model.layers.0.self_attn.o_proj.weight": torch.randn((2048, 2048)),
           "model.layers.0.self_attn.q_proj.weight": torch.randn((2048, 2048)),

ai_edge_torch/generative/test/test_experimental_ekv.py

@@ -16,20 +16,23 @@
 
 import unittest
 
-import numpy as np
-import torch
-
 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
 
 
 class TestExternalKVLayers(unittest.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)
+  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
+    )
     config = cfg.ModelConfig(
         kv_cache_max_len=kv_cache_max_len,
         embedding_dim=head_dim,
@@ -56,23 +59,31 @@ class TestExternalKVLayers(unittest.TestCase):
     entry = kv.caches[0]
     # single-slice update
     input_pos = torch.tensor([1])
-    k_slice = v_slice = torch.full((1, 1, NUM_QG, HEAD_DIM), 5, dtype=torch.float)
+    k_slice = v_slice = torch.full(
+        (1, 1, NUM_QG, HEAD_DIM), 5, dtype=torch.float
+    )
     updated_entry = kv_utils.update(entry, input_pos, k_slice, v_slice)
     self.assertEqual(
-        updated_entry.k_cache.numpy().flatten().tolist(), [0, 0, 5, 5, 0, 0, 0, 0]
+        updated_entry.k_cache.numpy().flatten().tolist(),
+        [0, 0, 5, 5, 0, 0, 0, 0],
     )
     self.assertEqual(
-        updated_entry.v_cache.numpy().flatten().tolist(), [0, 0, 5, 5, 0, 0, 0, 0]
+        updated_entry.v_cache.numpy().flatten().tolist(),
+        [0, 0, 5, 5, 0, 0, 0, 0],
     )
     # multi-slice update
     input_pos = torch.tensor([0, 3])
-    k_slice = v_slice = torch.full((1, 2, NUM_QG, HEAD_DIM), 7, dtype=torch.float)
+    k_slice = v_slice = torch.full(
+        (1, 2, NUM_QG, HEAD_DIM), 7, dtype=torch.float
+    )
     updated_entry = kv_utils.update(entry, input_pos, k_slice, v_slice)
     self.assertEqual(
-        updated_entry.k_cache.numpy().flatten().tolist(), [7, 7, 0, 0, 0, 0, 7, 7]
+        updated_entry.k_cache.numpy().flatten().tolist(),
+        [7, 7, 0, 0, 0, 0, 7, 7],
     )
     self.assertEqual(
-        updated_entry.v_cache.numpy().flatten().tolist(), [7, 7, 0, 0, 0, 0, 7, 7]
+        updated_entry.v_cache.numpy().flatten().tolist(),
+        [7, 7, 0, 0, 0, 0, 7, 7],
     )
 
   def test_serialization(self):

ai_edge_torch/generative/test/test_model_conversion.py

@@ -18,15 +18,14 @@ import os
 import tempfile
 import unittest
 
-import numpy as np
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.gemma import gemma
 from ai_edge_torch.generative.examples.phi2 import phi2
 from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache  # NOQA
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
 from ai_edge_torch.testing import model_coverage
+import numpy as np
+import torch
 
 
 class TestModelConversion(unittest.TestCase):