ai-edge-torch-nightly 0.3.0.dev20240909__py3-none-any.whl → 0.3.0.dev20240913__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_convolution.py

@@ -12,22 +12,171 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Provides lowering for coreaten to mlir stablehlo op: Convolution"""
+"""Provides lowering for coreaten to stablehlo for Convolution."""
 
 import math
 from typing import Optional
 
+from ai_edge_torch.odml_torch.lowerings import registry
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
 import torch
 
-from .registry import lower
+
+def make_padding(padding):
+  """Change the padding from pytorch to stablehlo style.
+
+  Stablehlo allows start and end padding for each dimension while aten only
+  allows symmetric padding and so only has one number per dimension.
+
+  Args:
+    padding: The padding of the convolution
+
+  Returns:
+    The padding in stablehlo style
+  """
+  return tuple((p, p) for p in padding)
+
+
+def create_conv_dimension_numbers(lhs, transposed: bool = False):
+  """Create the dimension numbers for the convolution.
+
+  Args:
+    lhs: The input tensor
+    transposed: Whether the convolution is transposed
+
+  Returns:
+    The dimension numbers for the convolution
+  """
+  num_spatial_dims = len(lhs.type.shape) - 2
+  spatial_dimensions = []
+  for i in range(0, num_spatial_dims):
+    spatial_dimensions.append(i + 2)
+
+  # Regular kernels are OIHW
+  # TransposedConv kernels are IOHW
+  dimension_numbers = stablehlo.ConvDimensionNumbers.get(
+      input_batch_dimension=0,
+      input_feature_dimension=1,
+      input_spatial_dimensions=spatial_dimensions,
+      kernel_input_feature_dimension=0 if transposed else 1,
+      kernel_output_feature_dimension=1 if transposed else 0,
+      kernel_spatial_dimensions=spatial_dimensions,
+      output_batch_dimension=0,
+      output_feature_dimension=1,
+      output_spatial_dimensions=spatial_dimensions,
+  )
+  return dimension_numbers
+
+
+def infer_output_shape(
+    lhs,
+    rhs,
+    stride,
+    dilation,
+    padding,
+    transposed: bool = False,
+    output_padding: list[int] = 0,
+):
+  """Infer the output shape of the convolution.
+
+  Args:
+    lhs: The input tensor
+    rhs: The kernel tensor
+    stride: The stride of the convolution (dilation of input in transposed conv)
+    dilation: The kernel dilation of the convolution
+    padding: The padding of the convolution
+    transposed: Whether the convolution is transposed
+    output_padding: The output padding of the convolution
+
+  Returns:
+    The output shape of the convolution
+  """
+  lhs_type: ir.RankedTensorType = lhs.type
+  lhs_shape: list[int] = lhs_type.shape
+  rhs_shape: list[int] = rhs.type.shape
+
+  # Input layout is: (N)CHW and Kernel layout is: (O)IHW for regular conv
+  # Input layout is: (N)CHW and Kernel layout is: I(O)HW for transposed conv
+  output_shape = (
+      [lhs_shape[0], rhs_shape[1]]
+      if transposed
+      else [lhs_shape[0], rhs_shape[0]]
+  )
+  num_spatial_dims = len(lhs.type.shape) - 2
+
+  # looping over the spatial dims (skipping the first 2 dims which are
+  # batch and features)
+  for spatial_dim in range(0, num_spatial_dims):
+    dim = spatial_dim + 2
+    dim_size = lhs_shape[dim]
+    kernel_dim_size = rhs_shape[dim]
+
+    if transposed:
+      output_dim_size = (
+          (dim_size - 1) * stride[spatial_dim]
+          - 2 * padding[spatial_dim]
+          + dilation[spatial_dim] * (kernel_dim_size - 1)
+          + output_padding[spatial_dim]
+          + 1
+      )
+    else:
+      output_dim_size = math.floor(
+          (
+              (
+                  dim_size
+                  + 2 * padding[spatial_dim]
+                  - dilation[spatial_dim] * (kernel_dim_size - 1)
+                  - 1
+              )
+              / stride[spatial_dim]
+          )
+          + 1
+      )
+
+    output_shape.append(output_dim_size)
+
+  return output_shape
+
+
+def build_transpose_conv(
+    lctx,
+    output_type: ir.RankedTensorType,
+    lhs: ir.Value,
+    rhs: ir.Value,
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    output_padding: list[int],
+    groups: int,
+):
+  lhs_type: ir.RankedTensorType = lhs.type
+  num_spatial_dims = len(lhs_type.shape) - 2
+  rhs = stablehlo.reverse(rhs, list(range(2, 2 + num_spatial_dims)))
+
+  kernel_size = rhs.type.shape
+  # We need to additional padding on the input to get the right output size.
+  adjusted_padding = [
+      dilation[dim] * (kernel_size[dim + 2] - 1) - padding[dim]
+      for dim in range(num_spatial_dims)
+  ]
+  return stablehlo.convolution(
+      result=output_type,
+      lhs=lhs,
+      rhs=rhs,
+      dimension_numbers=create_conv_dimension_numbers(lhs, True),
+      feature_group_count=groups,
+      batch_group_count=1,
+      padding=make_padding(adjusted_padding),
+      lhs_dilation=stride,
+      rhs_dilation=dilation,
+  )
 
 
 # convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride,
 #   SymInt[] padding, SymInt[] dilation, bool transposed,
 #   SymInt[] output_padding, SymInt groups) -> Tensor
-# @lower(torch.ops.aten.convolution)
+@registry.lower(torch.ops.aten.convolution)
 def _aten_convolution(
     lctx,
     lhs: ir.Value,
@@ -40,80 +189,53 @@ def _aten_convolution(
     output_padding: list[int],
     groups: int,
 ):
-  if transposed:
-    raise NotImplementedError("Transposed convolution is not implemented.")
 
-  if bias is not None:
-    raise NotImplementedError("Bias on convolution is not implemented.")
-
-  # Stablehlo allows start and end padding for each dimension while aten only
-  # allows symmetric padding and so only has one number per dimension.
-  def make_padding(padding):
-    return tuple((p, p) for p in padding)
-
-  def create_conv_dimension_numbers():
-    num_spatial_dims = len(lhs.type.shape) - 2
-    spatial_dimensions = []
-    for i in range(0, num_spatial_dims):
-      spatial_dimensions.append(i + 2)
-
-    dimension_numbers = stablehlo.ConvDimensionNumbers.get(
-        input_batch_dimension=0,
-        input_feature_dimension=1,
-        input_spatial_dimensions=spatial_dimensions,
-        kernel_input_feature_dimension=1,
-        kernel_output_feature_dimension=0,
-        kernel_spatial_dimensions=spatial_dimensions,
-        output_batch_dimension=0,
-        output_feature_dimension=1,
-        output_spatial_dimensions=spatial_dimensions,
+  # TODO(b/365559296) Add support for output_padding
+  if any(output_padding):
+    raise NotImplementedError(
+        "Output padding on convolution is not implemented."
     )
-    return dimension_numbers
-
-  def infer_output_shape():
-    lhs_type: ir.RankedTensorType = lhs.type
-    lhs_shape: list[int] = lhs_type.shape
-    rhs_shape: list[int] = rhs.type.shape
-
-    # Input layout is: (N)CHW and Kernel layout is: (O)IHW
-    output_shape = [lhs_shape[0], rhs_shape[0]]
-    num_spatial_dims = len(lhs.type.shape) - 2
-
-    # looping over the spatial dims (skipping the first 2 dims which are
-    # batch and features)
-    for spatial_dim in range(0, num_spatial_dims):
-      dim_size = lhs_shape[spatial_dim + 2]
-      kernel_dim_size = rhs_shape[spatial_dim + 2]
-
-      # for example, a dilation of 2 increases the dimension size by 2
-      dim_size *= dilation[spatial_dim]
-
-      # padding added to both sides
-      dim_size += 2 * padding[spatial_dim]
-
-      output_dim_size = math.ceil(
-          (dim_size - kernel_dim_size + 1) / stride[spatial_dim]
-      )
-
-      output_shape.append(output_dim_size)
-
-    return output_shape
 
   lhs_type: ir.RankedTensorType = lhs.type
-
-  op = stablehlo.ConvolutionOp(
-      result=ir.RankedTensorType.get(
-          infer_output_shape(), lhs_type.element_type
-      ),
-      lhs=lhs,
-      rhs=rhs,
-      dimension_numbers=create_conv_dimension_numbers(),
-      feature_group_count=groups,
-      batch_group_count=1,
-      window_strides=stride,
-      padding=make_padding(padding),
-      lhs_dilation=(1,) * len(stride),
-      rhs_dilation=dilation,
+  output_shape = infer_output_shape(
+      lhs, rhs, stride, dilation, padding, transposed, output_padding
+  )
+  output_type = ir.RankedTensorType.get(
+      output_shape,
+      lhs_type.element_type,
   )
 
-  return op.result
+  if transposed:
+    res = build_transpose_conv(
+        lctx,
+        output_type,
+        lhs,
+        rhs,
+        stride,
+        padding,
+        dilation,
+        output_padding,
+        groups,
+    )
+  else:
+    res = stablehlo.convolution(
+        result=output_type,
+        lhs=lhs,
+        rhs=rhs,
+        dimension_numbers=create_conv_dimension_numbers(lhs),
+        feature_group_count=groups,
+        batch_group_count=1,
+        window_strides=stride,
+        padding=make_padding(padding),
+        rhs_dilation=dilation,
+    )
+
+  if bias is not None:
+    # broadcast [C] to [NCHW]
+    broadcasted_bias = stablehlo.broadcast_in_dim(output_type, bias, [1])
+    res = stablehlo.add(
+        lhs=res,
+        rhs=broadcasted_bias,
+    )
+
+  return res

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -105,7 +105,6 @@ lower_by_torch_xla2(torch.ops.aten.clamp.default)
 lower_by_torch_xla2(torch.ops.aten.clone)
 lower_by_torch_xla2(torch.ops.aten.clone.default)
 lower_by_torch_xla2(torch.ops.aten.constant_pad_nd)
-lower_by_torch_xla2(torch.ops.aten.convolution)
 lower_by_torch_xla2(torch.ops.aten.cos)
 lower_by_torch_xla2(torch.ops.aten.cosh)
 lower_by_torch_xla2(torch.ops.aten.cumsum)
@@ -168,7 +167,6 @@ lower_by_torch_xla2(torch.ops.aten.mul.Scalar)
 lower_by_torch_xla2(torch.ops.aten.mul.Tensor)
 lower_by_torch_xla2(torch.ops.aten.native_batch_norm)
 lower_by_torch_xla2(torch.ops.aten.native_group_norm)
-lower_by_torch_xla2(torch.ops.aten.native_layer_norm)
 lower_by_torch_xla2(torch.ops.aten.native_layer_norm_backward)
 lower_by_torch_xla2(torch.ops.aten.ne)
 lower_by_torch_xla2(torch.ops.aten.neg)

ai_edge_torch/odml_torch/lowerings/_layer_norm.py

@@ -0,0 +1,78 @@
+# 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 lowering for coreaten to stablehlo for LayerNorm."""
+
+import math
+from typing import Optional
+from ai_edge_torch.odml_torch.lowerings import registry
+from ai_edge_torch.odml_torch.lowerings import utils
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+
+# native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight,
+# Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)
+@registry.lower(torch.ops.aten.native_layer_norm)
+def _aten_native_layer_norm(
+    lctx,
+    data: ir.Value,
+    normalized_shape: list[int],
+    weight: Optional[ir.Value],
+    bias: Optional[ir.Value],
+    eps: float,
+):
+  data_type: ir.RankedTensorType = data.type
+  unnormalized_count = math.prod(data_type.shape) // math.prod(normalized_shape)
+  dest_shape = [
+      1,
+      unnormalized_count,
+      math.prod(normalized_shape),
+  ]
+  dest_type = ir.RankedTensorType.get(dest_shape, data_type.element_type)
+
+  reshaped_data = stablehlo.reshape(dest_type, data)
+
+  one = utils.splat(1, data_type.element_type, [unnormalized_count])
+  zero = utils.splat(0, data_type.element_type, [unnormalized_count])
+  output, mean, var = stablehlo.batch_norm_training(
+      reshaped_data, one, zero, eps, 1
+  )
+  eps_splat = utils.splat(eps, var.type.element_type, var.type.shape)
+  rstd = stablehlo.rsqrt(stablehlo.add(var, eps_splat))
+
+  stats_shape = data_type.shape[: -1 * len(normalized_shape)] + [1] * len(
+      normalized_shape
+  )
+  stats_type = ir.RankedTensorType.get(stats_shape, data_type.element_type)
+  mean = stablehlo.reshape(stats_type, mean)
+  rstd = stablehlo.reshape(stats_type, rstd)
+
+  output = stablehlo.reshape(data_type, output)
+
+  data_rank = len(data_type.shape)
+  normalized_rank = len(normalized_shape)
+  if weight is not None:
+    weight = stablehlo.broadcast_in_dim(
+        data_type, weight, list(range(data_rank - normalized_rank, data_rank))
+    )
+    output = stablehlo.multiply(weight, output)
+  if bias is not None:
+    bias = stablehlo.broadcast_in_dim(
+        data_type, bias, list(range(data_rank - normalized_rank, data_rank))
+    )
+    output = stablehlo.add(bias, output)
+
+  return output, mean, rstd

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20240909"
+__version__ = "0.3.0.dev20240913"

{ai_edge_torch_nightly-0.3.0.dev20240909.dist-info → ai_edge_torch_nightly-0.3.0.dev20240913.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240909
+Version: 0.3.0.dev20240913
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.3.0.dev20240909.dist-info → ai_edge_torch_nightly-0.3.0.dev20240913.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,116
 ai_edge_torch/config.py,sha256=PCd9PVrbUNeVIUDFUCnW4goDWU4bjouK28yMYU6VOi0,877
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=NYV6Mkaje_ditIEI_s_7nLP_-8i4kbGM8nRzieVkbUI,5397
-ai_edge_torch/version.py,sha256=r0y6crIySNGhJqtljkzyHxb1XMvLji2VLajLfUjW8b4,706
+ai_edge_torch/version.py,sha256=2_ahYhvytovu9mWRifMKeqx6-0JbD7-iV5FXU890d7Y,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=kcv_QgNgeyDmrqwdzHicGNP68w6zF7GJg7YkMEIXp4Q,3759
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -26,7 +26,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=pUYSXuqFg8CAeJ8JkoYf7S0RDLRPVuZUwVOd0xObM6w,14411
+ai_edge_torch/_convert/test/test_convert.py,sha256=FSufFZEeTLBpUnzE1Iy-LvNN0mhDynWMNg7Mei8RpLQ,14973
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -39,27 +39,20 @@ ai_edge_torch/debug/test/test_search_model.py,sha256=-RuU0QsjqkfzZF2IbeA55MoeVOa
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/experimental/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/experimental/gemma/convert_to_tflite.py,sha256=lpiPFSh3SJd6WwuZ0QegSva3__iSz2tUD7L7QfkAe4I,3085
-ai_edge_torch/generative/examples/experimental/gemma/gemma.py,sha256=aCoD86pf4nuquUMk7MOR-jsN5FqvySSEuMx9Psxjblk,7261
-ai_edge_torch/generative/examples/experimental/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/experimental/phi/convert_to_tflite.py,sha256=DavrdGmqUgoThsGNRv3LXMW5tvJdYEvj66Hf1XRqkXU,3055
-ai_edge_torch/generative/examples/experimental/phi/phi2.py,sha256=Jxf3ZyYDpS78l6uh4_LGGIcHawrOhZ1vHoHFVxRaK40,6789
-ai_edge_torch/generative/examples/experimental/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/experimental/tiny_llama/convert_to_tflite.py,sha256=xPVvHQjLJHFiRv_-Fy2sDm0Aft7SG8SXiV6o3rF03cQ,3108
-ai_edge_torch/generative/examples/experimental/tiny_llama/tiny_llama.py,sha256=nUm0SQbCTmNAc5u-C9gbQRFPt7GDvUt6UjH6doTvH-I,6817
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=pseJExH35lSAK0ZtzSHB1sFtRtF_EuT2xcSpGU0gKVI,2524
-ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=w589IJETATd6Z9_1XCIWbrlCV3E92X_5ac3VVCVFXG0,2522
-ai_edge_torch/generative/examples/gemma/gemma.py,sha256=lc1-CfIObHj9D5VJy78BOtGTrQM4TYMI6NfVi8KM5qA,6747
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=OcUQLFR136e3QRVXRnmtYnRHXyHJS9EYEFlJ1ymXyRY,8859
-ai_edge_torch/generative/examples/phi2/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=ON6zLO-nFS8eJ2yhyWzT5x2Somr-Ca-VjpjT7OGFU10,2506
-ai_edge_torch/generative/examples/phi2/phi2.py,sha256=FFnhv1kx4fHRhSeOreLGj8kAqPnmkz9pD1RRSDVlM_w,6332
+ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=ZJvw8uFVu7FEJ7eXfpzn-pPKgPELoxkGz4Zg7LKKMSI,3048
+ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=hM-fwjZG53p1UE_lkovLMmHRDHleJsb6_0ib0_k0v54,3040
+ai_edge_torch/generative/examples/gemma/gemma.py,sha256=uejk9Mi85uRuFYIUi5XI58rf4K7TFeE5cZ1flejF8EE,7473
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=H0scyAdqRyV2wwaFx1LAa3A5oYn1C5tTdPWvbDTd_SQ,10256
+ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=vqEpZVmB0_wMKcAl6RXm7W57DqPTzEdVVN6W2Z-QYzI,3011
+ai_edge_torch/generative/examples/phi/phi2.py,sha256=wjTLCfCUDcLqvVsrPH-Wx04pOKeuigZCWHO3gL1WOEA,7072
+ai_edge_torch/generative/examples/smallm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/generative/examples/smallm/convert_to_tflite.py,sha256=aqqxQMBBO_dtGB1iZ1tpF8hbGpdZkx0VIz62ZqfVMCc,3036
+ai_edge_torch/generative/examples/smallm/smallm.py,sha256=mzlbXxCCB10FN03QDRoPXw-cbucQM_O_Hs8hqLZAvck,4002
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
-ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=0WniBWQ6_NcQc5WycX3YRRX7Os9AGQSxfc1m2HKBqg8,4479
+ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=evl5Rn_Hlp9-BsNmcf6liXa2syET3-Fz-zVaWjqPKx8,4657
 ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=7ra36nM5tQwSw-vi6QCFLx5IssZhT-6yVK4H3XsAc4w,5044
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=slieF2-QcDCwd4DRZ7snsZIphT97IXpp4plRRsRSwL8,13983
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=7oUIJ6HO0vmlhFdkXpqGm9KTB-eM4Ob9VrHSDlIGFOg,30926
@@ -74,29 +67,28 @@ ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py,sha256=ZE6H
 ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py,sha256=RxR5rw0wFFm_5CfAY-3-EIz83vhM9EKye8Bb5zBb0Ok,1341
 ai_edge_torch/generative/examples/t5/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=CZVuNEL8OHPkdsz70WOvNpTJ9LFkiDnlwgJiXfUZCVk,4548
-ai_edge_torch/generative/examples/t5/t5.py,sha256=Zobw5BV-PC0nlU9Z6fzb2O07rMeU8vGIk-KtKp9D_H0,20871
-ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=1lvbSlzyBwmd5Bs7-Up_v4iJQkCPIJx2RmMkLgy7l2Q,8508
+ai_edge_torch/generative/examples/t5/t5.py,sha256=Ekg92OwIXSkSRii9OY-mp3-SExtsxOdoIDTFxm25hso,21304
+ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=l01oYyJo77INzRwN4xqXquaFQPvCFBFF5zOnmGVb3Hg,8731
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=5wj2RmQRIwD6O_R_pp-A_7gKGSdHWDSXyis97r1ELVI,5622
-ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py,sha256=l9swUKTcDtnTibNSNExaMgLvDeJ4Er2tVh5ZW1EtRgk,5809
-ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=mQkcpSe6HlRLMkIRCEHc9ZXL7jxEp9RWSGUQjjd-r2w,4841
+ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=QyLeCqDnk71WvvFH68g9UeF-HytonSk1ItGF9dc7Zj8,5854
+ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=oX_D_kU9PegBX3Fx9z_J3a1Oh2PF05F0nwZNxyLgQNA,5880
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=CLRqO7ycMbpy7J3_Czp1sLx6hcdwGD9zVq04yRba0e8,2550
-ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=4ku0ni3MOWamhPrzLap0BmtdNFk7CH0hwjPNoRAKpvQ,6278
+ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=y4LiWhwgflqrg4WWh3wq5ei3VOT_cV0A62x62qptQiM,3070
+ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=Mnn_aMImR1CpC_T0CMKlp3XgoLyR7N56VR3blVSnMHQ,7007
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=fmNNXawJ722M4cTUuTx289rT0NHxBEsOy_k8baqCOms,1173
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=sXis0U4u-RoIp_NyrmWJNnqFqpqRuZOrhfsJIO6rMps,2028
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=2UujQePRJ1LK02PN-hGcuMu0ooCJC6ETfPvzEYVFyho,12284
+ai_edge_torch/generative/layers/attention.py,sha256=d9yLaqxPCtClhNUmauOEFBKxhLnsXdN3NiYy1WspIPI,12826
 ai_edge_torch/generative/layers/attention_utils.py,sha256=68GXGR2HSWBFViTxX7cHifzVG-kcLS2IL2tQJPIpupg,7344
-ai_edge_torch/generative/layers/builder.py,sha256=xb7rjADv3Jm4qfmlYtg6oLLe7ReDE9UjsEqiejPpDD8,4346
+ai_edge_torch/generative/layers/builder.py,sha256=6jDNaa_djF32AjxIJtaDGBzlj3zlvl1yZivK3gC4j94,4424
 ai_edge_torch/generative/layers/feed_forward.py,sha256=uto7xtwx6jPkk1GZ2x7pSTentQzRrPSKw4_PSE12ahA,3525
-ai_edge_torch/generative/layers/kv_cache.py,sha256=Ob8QeXWW5xt-6hcGA0uoC48eRQ8lfvKca8JbWtFx2CE,3082
-ai_edge_torch/generative/layers/model_config.py,sha256=WpZ9djUBAZddyeSODHDaVMG37EQqfzGGrlMPi8AA-Hc,5752
-ai_edge_torch/generative/layers/normalization.py,sha256=u8lv0p-ktKcRqCDlOqZQa9WQcfDK9JM2IaUQFQdn7xs,1860
+ai_edge_torch/generative/layers/kv_cache.py,sha256=FveTTO0z_yi0-ZdGMuamzSvuInn6B4lesKZ4PHT2Vmg,6088
+ai_edge_torch/generative/layers/model_config.py,sha256=mil4RkGuNFBDKo3gPd9QnfGKLKPZWX9Gz2_q9hX8sNU,6407
+ai_edge_torch/generative/layers/normalization.py,sha256=iod9oNkoDS5m-yFY_Y_XMyvCU5a88ESd_s5WY34ErKA,6129
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=CZqOoibLcHvUgrgaIIWAlmk3XgE2inzx340MN-npLoU,1347
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=VW-VP8e7FTSPCdu-6DVxpwNrIdgX0R_kq6F6MSEiyXE,3848
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=V4zUAqjWeBseMPG9B-93LDv1LM3Dds6Q-H0NxY0koSA,27212
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=cpygyJccLq6KHKxV7oz4YKh529YLjC9isupnsVmPi0A,27190
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=NvBJj09a7ZC-ChGE_ex-_kLnE_fjzrY6txbLSh1pMKA,9208
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -107,15 +99,16 @@ ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC
 ai_edge_torch/generative/quantize/quant_recipes.py,sha256=0Kvr_o7pbMnE8VMe6Ml0FBxkHM6RJ3C14B2I1mjItjc,2030
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVufetPvkDoD7rUowIOSKV9oV5Kk,1418
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/test/test_experimental_ekv.py,sha256=8qv_eVtJW9GPvBEf2hPQe3tpdJ33XShya6MCX1FqrZM,4355
-ai_edge_torch/generative/test/test_loader.py,sha256=_y5EHGgoNOmCuYonsB81UJScHVsTAQXUVd44czMAw6k,3379
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=b3InJ8Rx03YtHpE9h-j0pSXAY1cCf-dLlx4Y5LSJnRQ,5174
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=9JXcd-rX8MpsYeEWUFEXf783GOwYOLY64KzDfFdmRJ8,4484
+ai_edge_torch/generative/test/test_kv_cache.py,sha256=W6Bh0gYDzmwb0j9HdD5_D7Z7FPToP2HSyFrmwIXuFqo,3793
+ai_edge_torch/generative/test/test_loader.py,sha256=8y74ChO3CZCfEi1eCf3-w47kRgAI4qPYCXpi8rTQXMA,3378
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=SIv7_sc5qHvbHFN8SbAfY00iXGvH7J6cJLkERU_cd5k,5888
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=F3q3K9ZgWBzlLy4WpE8-w6UWSuJ-UoJwMm3N6Zb3Y14,5016
 ai_edge_torch/generative/test/test_quantize.py,sha256=kY_NRpF-v1i4clqI1CFFWEagJv-5PzBDkeJ2fInl9_w,5913
+ai_edge_torch/generative/test/utils.py,sha256=YvEhO2HIj1LkBs5du1UxY-cGRW9HMyAYsOUhgsTrTpA,1796
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/loader.py,sha256=6J0aAP6-6LySeqeYIHKcchr5T9cVtSO34aoDr3V9gxY,12726
+ai_edge_torch/generative/utilities/loader.py,sha256=kn4TCgGAG8s4mdvPITimOBCaVyn04Ksz4gZIleFYF1o,12754
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=pKp3AMSbS3otCvgwJRF5M1l4JRNKk-aCKimXzIMSrds,35679
-ai_edge_torch/generative/utilities/t5_loader.py,sha256=_UXcc1QKT-S92hikfo-fTBFhnYLzROqcyRqKonVsqj4,16885
+ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
@@ -145,11 +138,12 @@ ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkgH7x38qPlwUUpD7S15Q,730
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=drN3L0uTsSjkluKgt6Ngq7b5HLReE_7iAitHpZ9PKqE,5428
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
-ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=GqYk6oBJw7KWeG4_6gxSu_OvYhjJcC2FpGzWPPEdH6w,933
+ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=dE_qzh-OnCNjWzqs1-PHs5PNlRF726qMQKM3tkwAzEs,959
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=wV8AUK8dvjLUy3qjqw_IxpiYVDWUMPNZRfi3XYE_hDs,6972
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
-ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=B6BILeu-UlwGB1O6g7111X1TaIFznsfxXrB72ygBsBA,3885
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=I0Y4IK7Zap8m6xfxMw7DfQ9Mg4htKOoypdHVAMHqx9c,10669
+ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=v1VdKmL8YLJv3PR9VgyNghO83A25PpTzY2ZUAJqlq3Q,6847
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=Ii1akrKLhRTkZ715JxXBBGKv3jGfXReXMQCYNzSnxmM,10567
+ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=1ePJs7oIdUkVdMddFsXMc53qTkEKqGz0ZhQQoNzBa10,2862
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=ES3x_RJ22T5rlmMrlomex2DdcZbhlyVJ7_HS3rjz3Uk,2851
 ai_edge_torch/odml_torch/lowerings/utils.py,sha256=NczqpsSd3Fn7yVcPC3qllemiZxxDAZgcW1T5l8-W9fE,5593
@@ -161,8 +155,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20240909.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240909.dist-info/METADATA,sha256=s7SAIUvFciy8peNKMHvyhoNQWYx67Jerz4foeV7KiE0,1859
-ai_edge_torch_nightly-0.3.0.dev20240909.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240909.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240909.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240913.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240913.dist-info/METADATA,sha256=ahbsMN1e0Tuq_LmrkB6NE-VgVTC65KEiZX3VVmTbcWQ,1859
+ai_edge_torch_nightly-0.3.0.dev20240913.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240913.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240913.dist-info/RECORD,,