ai-edge-torch-nightly 0.3.0.dev20250122__py3-none-any.whl → 0.3.0.dev20250124__py3-none-any.whl

ai_edge_torch/_config.py

@@ -65,5 +65,14 @@ class _Config:
   def enable_group_norm_composite(self, value: bool):
     os.environ["ENABLE_GROUP_NORM_COMPOSITE"] = "y" if value else "n"
 
+  @property
+  def layout_optimize_partitioner(self) -> str:
+    """The algorithm to use for layout optimization."""
+    return os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", "DEFAULT")
+
+  @layout_optimize_partitioner.setter
+  def layout_optimize_partitioner(self, value: str):
+    os.environ["AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER"] = str(value).upper()
+
 
 config = _Config()

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -201,8 +201,14 @@ def _aten_group_norm_checker(node):
   return NHWCable(can_be=can_be, must_be=must_be)
 
 
-@nhwcable_node_checkers.register(aten.native_group_norm)
+@nhwcable_node_checkers.register(aten.native_group_norm.default)
 def _aten_native_group_norm_checker(node):
+  # aten.group_norm is removed from the decomp table, so aten.native_group_norm
+  # should never exist in the graph. However, torch 2.5.1 could ignore the
+  # decomp table updates, so still add this native_group_norm checker and
+  # rewriter to be safe.
+  # The checker and rewriter are the same as the ones for aten.group_norm.
+
   val = node.meta.get("val")
   if (
       not isinstance(val, (list, tuple))
@@ -210,13 +216,10 @@ def _aten_native_group_norm_checker(node):
       or not hasattr(val[0], "shape")
   ):
     return NHWCable(can_be=False, must_be=False)
-  if len(node.args) >= 3 and (
-      node.args[1] is not None or node.args[2] is not None
-  ):
-    # Disable NHWC rewriter due to precision issue with weight and bias.
-    # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
-    return NHWCable(can_be=False, must_be=False)
-  return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
+
+  can_be = len(val[0].shape) == 4
+  must_be = can_be and ai_edge_torch.config.enable_group_norm_composite
+  return NHWCable(can_be=can_be, must_be=must_be)
 
 
 # ==== Ops must be NCHW

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -391,34 +391,32 @@ def _aten_native_group_norm(node):
       eps: float,
       **kwargs,
   ):
-    input_reshaped = torch.reshape(
-        input,
-        [
-            batch_size,
-            flattened_inner_size,
-            num_groups,
-            num_channels // num_groups,
-        ],
-    )
-    reduction_dims = [1, 3]
-
-    biased_var, mean = torch.var_mean(
-        input_reshaped, dim=reduction_dims, unbiased=False, keepdim=True
+    is_composite_supported = (
+        ai_edge_torch.config.enable_group_norm_composite
+        and weight is not None
+        and bias is not None
     )
-    rstd = torch.rsqrt(biased_var + eps)
-
-    out = (input_reshaped - mean) * rstd
-    out = torch.reshape(out, input.shape)
 
-    if weight is not None:
-      out = out * weight
-    if bias is not None:
-      out = out + bias
+    builder = None
+    if is_composite_supported:
+      builder = StableHLOCompositeBuilder(
+          name="odml.group_norm",
+          attr={
+              "num_groups": num_groups,
+              "epsilon": eps,
+              "reduction_axes": [3],
+              "channel_axis": 3,
+          },
+      )
+      input, weight, bias = builder.mark_inputs(input, weight, bias)
 
-    mean = torch.squeeze(mean, reduction_dims)
-    rstd = torch.squeeze(rstd, reduction_dims)
+    input = utils.tensor_to_nchw(input)
+    output = aten.group_norm.default(input, num_groups, weight, bias, eps=eps)
+    output = utils.tensor_to_nhwc(output)
 
-    return out, mean, rstd
+    if builder is not None:
+      output = builder.mark_outputs(output)
+    return (output, None, None)
 
   node.target = native_group_norm
 

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -18,6 +18,7 @@ import operator
 import os
 from typing import Union
 
+import ai_edge_torch
 from ai_edge_torch import fx_infra
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
@@ -261,10 +262,8 @@ class OptimizeLayoutTransposesPass(fx_infra.ExportedProgramPassBase):
     self.mark_const_nodes(exported_program)
 
     graph_module = exported_program.graph_module
-    partitioner = os.environ.get(
-        "AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None
-    )
-    if partitioner == "MINCUT":
+    partitioner = ai_edge_torch.config.layout_optimize_partitioner
+    if partitioner in ("MINCUT", "OPTIMAL"):
       graph_module = layout_partitioners.min_cut.partition(graph_module)
     elif partitioner == "GREEDY":
       graph_module = layout_partitioners.greedy.partition(graph_module)

ai_edge_torch/generative/layers/kv_cache.py

@@ -81,7 +81,8 @@ class KVCache:
     """
     caches = [
         KVCacheEntry.from_model_config(
-            config.kv_cache_max,
+            config.kv_cache_max if not config.block_config(idx).kv_cache_max_len 
+            else config.block_config(idx).kv_cache_max_len,
             config.block_config(idx).attn_config,
             dtype,
             device,

ai_edge_torch/generative/layers/model_config.py

@@ -164,6 +164,9 @@ class TransformerBlockConfig:
   parallel_residual: bool = False
   # The Attention computation will include relative positional bias.
   relative_attention: bool = False
+  # KV Cache length for this block. Only used when attention types are different
+  # across blocks
+  kv_cache_max_len: Optional[int] = None
 
 
 @dataclasses.dataclass
@@ -200,7 +203,8 @@ class ModelConfig:
   embedding_use_bias: bool = False
   # Image embedding parameters.
   image_embedding: Optional[ImageEmbeddingConfig] = None
-
+  # Number of image tokens 
+  num_mm_tokens_per_image: Optional[int] = None
   # Use bias term within LLM's HEAD.
   lm_head_use_bias: bool = False
   # Whether LLM's HEAD shares the weight of the embedding.

ai_edge_torch/generative/utilities/bmm_4d.py

@@ -0,0 +1,76 @@
+# 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.
+# ==============================================================================
+# Common utility functions for data loading etc.
+from dataclasses import dataclass
+import glob
+import os
+from typing import Sequence
+from ai_edge_torch.odml_torch import lowerings
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+
+# Use torch.library.custom_op to define a new custom operator.
+@torch.library.custom_op("ai_edge_torch::bmm_4d", mutates_args=())
+def bmm_4d(
+    lhs: torch.Tensor,
+    rhs: torch.Tensor,
+) -> torch.Tensor:
+  if not (lhs.ndim == 4 and rhs.ndim == 4):
+    raise ValueError("bmm_4d requires LHS and RHS have rank 4.")
+  d0_can_bcast = lhs.shape[0] == rhs.shape[0] or lhs.shape[0] == 1 or rhs.shape[0] == 1
+  d1_can_bcast = lhs.shape[1] == rhs.shape[1] or lhs.shape[1] == 1 or rhs.shape[1] == 1
+  if not (d0_can_bcast and d1_can_bcast):
+    raise ValueError("bmm_4d requires that dimensions 0 and 1 can broadcast.")
+
+  if not lhs.shape[-1] == rhs.shape[-1]:
+    raise ValueError("bmm_4d requires LHS and RHS have same last dimension.")
+
+  return torch.einsum("abcd,abed->abce", lhs, rhs)
+
+
+# Use register_fake to add a ``FakeTensor`` kernel for the operator
+@bmm_4d.register_fake
+def _(lhs, rhs):
+  return torch.einsum("abcd,abed->abce", lhs, rhs)
+
+
+@lowerings.lower(torch.ops.ai_edge_torch.bmm_4d)
+def _bmm_4d_lower(
+    lctx,
+    lhs: ir.Value,
+    rhs: ir.Value,
+):
+  dot_dnums = stablehlo.DotDimensionNumbers.get(
+      lhs_batching_dimensions=[0, 1],
+      rhs_batching_dimensions=[0, 1],
+      lhs_contracting_dimensions=(3,),
+      rhs_contracting_dimensions=(3,),
+  )
+  return stablehlo.dot_general(
+      ir.RankedTensorType.get(
+          (
+              lhs.type.shape[0],
+              lhs.type.shape[1],
+              lhs.type.shape[2],
+              rhs.type.shape[2],
+          ),
+          lhs.type.element_type,
+      ),
+      lhs,
+      rhs,
+      dot_dnums,
+  )

ai_edge_torch/generative/utilities/converter.py

@@ -174,6 +174,9 @@ def _export_helper(
           'input_pos': prefill_input_pos,
           'kv_cache': kv,
       }
+      if export_config.prefill_mask is not None:
+        sample_kwargs['mask'] = export_config.prefill_mask
+
       if lora is not None:
         prefill_signature_name += f'_lora_r{lora.get_rank()}'
         sample_kwargs['lora'] = lora
@@ -199,6 +202,8 @@ def _export_helper(
         'input_pos': decode_input_pos,
         'kv_cache': kv,
     }
+    if export_config.decode_mask is not None:
+      sample_kwargs['mask'] = export_config.decode_mask
     if lora is not None:
       sample_kwargs['lora'] = lora
 

ai_edge_torch/generative/utilities/model_builder.py

@@ -55,6 +55,9 @@ class ExportConfig:
   # On prefill signatures, should the model produce logit output?
   # When False, only decode signatures will produce output.
   output_logits_on_prefill: bool = False
+  # Attention masks given as inputs to the model.
+  prefill_mask: Optional[torch.Tensor] = None
+  decode_mask: Optional[torch.Tensor] = None
 
 
 class DecoderOnlyModel(nn.Module):

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -29,7 +29,7 @@ LoweringContext = context.LoweringContext
 
 @functools.cache
 def _log_usage(op):
-  logging.warning("Use jax lowering: %s", str(op))
+  logging.info("Use JAX lowering: %s", str(op))
 
 
 def lower_by_jax(op, ir_input_names=None):

ai_edge_torch/quantize/pt2e_quantizer_utils.py

@@ -21,8 +21,6 @@ from typing import Callable, Dict, List, NamedTuple, Optional
 import torch
 from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix
 from torch.ao.quantization.pt2e.graph_utils import find_sequential_partitions
-from torch.ao.quantization.pt2e.utils import _conv1d_bn_example_inputs
-from torch.ao.quantization.pt2e.utils import _conv2d_bn_example_inputs
 from torch.ao.quantization.pt2e.utils import _get_aten_graph_module_for_pattern
 from torch.ao.quantization.quantizer import QuantizationAnnotation
 from torch.ao.quantization.quantizer import QuantizationSpec
@@ -47,6 +45,28 @@ __all__ = [
     "propagate_annotation",
 ]
 
+# Example inputs for conv-bn1d patterns
+_conv1d_bn_example_inputs = (
+    torch.randn(1, 1, 3),  # x
+    torch.randn(1, 1, 1),  # conv_weight
+    torch.randn(1),  # conv_bias
+    torch.randn(1),  # bn_weight
+    torch.randn(1),  # bn_bias
+    torch.randn(1),  # bn_running_mean
+    torch.randn(1),  # bn_running_var
+)
+
+# Example inputs for conv-bn2d patterns
+_conv2d_bn_example_inputs = (
+    torch.randn(1, 1, 3, 3),  # x
+    torch.randn(1, 1, 1, 1),  # conv_weight
+    torch.randn(1),  # conv_bias
+    torch.randn(1),  # bn_weight
+    torch.randn(1),  # bn_bias
+    torch.randn(1),  # bn_running_mean
+    torch.randn(1),  # bn_running_var
+)
+
 
 @dataclass(eq=True, frozen=True)
 class QuantizationConfig:

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20250122.dist-info → ai_edge_torch_nightly-0.3.0.dev20250124.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250122
+Version: 0.3.0.dev20250124
 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.dev20250122.dist-info → ai_edge_torch_nightly-0.3.0.dev20250124.dist-info}/RECORD

@@ -1,8 +1,8 @@
 ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,1208
-ai_edge_torch/_config.py,sha256=PKtOtBOup-cM0wBdQxby6HzuhLhIC3oq-TBG8FF4znE,2161
+ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=hTwjlLVIp826evZRXkCm4XtYo3T7qGsjqbgW_Llf138,706
+ai_edge_torch/version.py,sha256=xloGd_dX0MD8k-quT07WLlEN1zIGVtCKu6xBSvjofrc,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -16,11 +16,11 @@ ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=Z6E3U7SYZv
 ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=HCOkj0k3NhaYbtfjE8HDXVmYhZ9fL5V_u6VunVh9mN4,2116
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=UKC-wM93-oe8spxyFqgybJ0TwnSRw8f-SOA2glCh2FA,890
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/_decomp_registry.py,sha256=aWO_zHDF4j_hokoKJQNFIFmua4ysXztsgS6pcyBUht0,1082
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=S_Bniv6jY16oOoFUzlyECQ0I2HDjG2D1MOI-QYPk3jQ,8061
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=7yEKSfXskXUk4tsd7c8vL155O-iU4eUjXCU5RSZHrbw,8204
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=4RyGUwR22bZqkn_TnptenFJodc_Q43f4_SBG7gmTbos,1621
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=zoAZ2TXKvxUnWnT11U4tx2uF0J5kkNXydgaW7JzfkXI,13811
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=OCFcPP618zH8IE12KTBQm2hRTtsaSeO3egvlOBUpNxA,13911
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=OhisegHY2j4cv_m9auCh9Mq9qmm1lUqpFLVO9X-oBlc,1032
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=lgoH32l6zAbWTCpa_4-RWkHjqbNaPsBnhSObLIX8dL4,10551
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=mr0MiLbaQmU-3S3KT-vb58kRWbNT3VJiCKY-K7_3tFg,10556
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/utils.py,sha256=YLMttMg5PdvXTtQ8lxpKb434UGVvYVALV1-xeuH4UGc,2131
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py,sha256=D8VX8SbCzfoyvPgMFHK7yxD7R-bzLxp2gfdKxgrWekA,742
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
@@ -128,9 +128,9 @@ ai_edge_torch/generative/layers/attention.py,sha256=GrAy8CT1pEsgRoB8JQP6PlnNYk8k
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=DhHIggaOQ2IAY4aRuMAuCLWZv1dBz5PYtmOEjkx9EQY,6291
+ai_edge_torch/generative/layers/kv_cache.py,sha256=sGGAZD0mWYuO4FukZfDbHXoxpBOBE9lTYICvZzDj5F8,6400
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
-ai_edge_torch/generative/layers/model_config.py,sha256=9yPEmWNw3-_2wXBmPmZ7RUKcPXHF2ZbJwksyQoXTA6M,7784
+ai_edge_torch/generative/layers/model_config.py,sha256=ZVRWEGw1BnLbLCuoR71kWGqQteKp-UM1YvMbbWYlkNw,7999
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
@@ -155,10 +155,11 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=yzMgXkiZxHUF
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=yNIZ-O6RdXYl8yuWM_sTENRxozPnKGS-TZRhiiTaraE,7515
+ai_edge_torch/generative/utilities/bmm_4d.py,sha256=2BMOYiFVUsl-bjxmLkrX4N7kpO0CnhB7eDYxm_iBCr8,2533
+ai_edge_torch/generative/utilities/converter.py,sha256=QIYxT-zATMzsD3LG-keRkxpJqDKXkbil4Se1KXthWFg,7726
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
-ai_edge_torch/generative/utilities/model_builder.py,sha256=3CQLxJ02pFIo2DlS-RCn9cT6OvR4NiIuYRH597UXLiI,6530
+ai_edge_torch/generative/utilities/model_builder.py,sha256=aXigoFEMLAKk7HQuWJM5ILs3igA4z2VH64ZCzCuBhDE,6671
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
@@ -197,7 +198,7 @@ ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=8mZTp_ybcMO3tDRQdlDP68BVeTw5
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=VhmeGFnB5hrUsALiVWV96JJOqPDrTIWouHjTvLuT5eU,2477
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=CJHWkmY4aAVQ5dmFsVc3Ox9TPkoLSNOfa96psD4CLRo,11561
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=fte81SZxgxeMcI3wWVKSTnUjIxVVilOJ6H3TybXyDmQ,11558
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
@@ -207,13 +208,13 @@ ai_edge_torch/odml_torch/lowerings/utils.py,sha256=pqM6mumpviFDHRaabp93CUAngzEZm
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
 ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV_TmSqEBPjes,15741
-ai_edge_torch/quantize/pt2e_quantizer_utils.py,sha256=eARD1LxLi5m7Z0n_psAkeX_AtUp4fNkE--oECBfivv4,36208
+ai_edge_torch/quantize/pt2e_quantizer_utils.py,sha256=nuO3w9gOj9sKcsTBBexVDw8UZnd06KsjNrFr_gyNaiA,36710
 ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9PphCRdO8o,3172
 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.dev20250122.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250122.dist-info/METADATA,sha256=-gaDJq2ogHs2JdgtvtQwpsXxZQkKB9e3yOLcC1XYz2I,1966
-ai_edge_torch_nightly-0.3.0.dev20250122.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250122.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250122.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/METADATA,sha256=WKVqBXJtXvMv3JfqtYKcl1GFgKrKtSbZ-tJAol5PPHk,1966
+ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/RECORD,,