ai-edge-torch-nightly 0.3.0.dev20241120__py3-none-any.whl → 0.3.0.dev20241122__py3-none-any.whl

ai_edge_torch/_convert/test/test_convert_composites.py

@@ -39,7 +39,6 @@ def _func_to_torch_module(func: Callable[..., torch.Tensor]):
   return TestModule(func).eval()
 
 
-@googletest.skip('Temporary outage due to changes for b/377531086')
 class TestConvertComposites(googletest.TestCase):
   """Tests conversion modules that are meant to be wrapped as composites."""
 

ai_edge_torch/generative/layers/rotary_position_embedding.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 # Implementation for Rotary Position embedding. https://arxiv.org/pdf/2104.09864.pdf
+from typing import Tuple
 import torch
 
 
@@ -36,3 +37,52 @@ def apply_rope(
   rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
   roped = (x * cos) + (rotated * sin)
   return roped.transpose(1, 2).type_as(x)
+
+
+def apply_rope_inline(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    input_pos: torch.Tensor,
+    n_elem: int,
+    base: int = 10_000,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Computes rotary positional embedding inline for a query and key.
+
+  Args:
+    q: the query tensor.
+    k: the key tensor.
+    input_pos: the sequence indices for the query and key
+    n_elem: number of elements of the head dimension for RoPE computation
+
+  Returns:
+    output the RoPE'd query and key.
+  """
+
+  if n_elem <= 0:
+    return q, k
+
+  theta = 1.0 / (base ** (torch.arange(0, n_elem, 2).float() / n_elem))
+  freq_exponents = (2.0 / n_elem) * torch.arange(
+      q.shape[-1] // 2, dtype=torch.float32
+  )
+  timescale = float(base) ** freq_exponents
+  radians = input_pos.clone().unsqueeze(0).unsqueeze(-1) / timescale.unsqueeze(
+      0
+  ).unsqueeze(0)
+  cos = torch.cos(radians).type_as(q)
+  sin = torch.sin(radians).type_as(q)
+
+  def apply(x, sin, cos):
+    x = x.transpose(1, 2)
+    b, h, s, d = x.shape
+    ans = torch.split(x, d // 2, dim=-1)
+    x1, x2 = ans
+    left = x1 * cos - x2 * sin
+    right = x2 * cos + x1 * sin
+    res = torch.cat([left, right], dim=-1)
+    res = res.transpose(1, 2)
+    return res
+
+  q_roped = apply(q, sin, cos)
+  k_roped = apply(k, sin, cos)
+  return q_roped, k_roped

ai_edge_torch/generative/test/test_custom_dus.py

@@ -0,0 +1,107 @@
+# 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.
+# ==============================================================================
+
+"""A suite of tests to validate the Dynamic Update Slice Custom Op."""
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import torch
+from torch import nn
+
+from absl.testing import absltest as googletest, parameterized
+
+
+def updated_slice_matches(buffer, update, index):
+  indexer = [slice(i, i + d) for i, d in zip(index, update.shape)]
+  buf = buffer[indexer]
+  return torch.allclose(buf, update)
+
+
+def intT(x):
+  return torch.tensor(x).int()
+
+
+class DUSMod(nn.Module):
+
+  def forward(self, buffer, update, index):
+    out = dynamic_update_slice(buffer, update, index)
+    out = out * 2
+    return out
+
+
+@googletest.skip('Enable this when odml_torch is default b/373387583')
+class TestCustomDUS(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      (
+          'DUS_whole_buffer',
+          torch.randn(1, 1280, 4, 64),
+          torch.randn([1, 1024, 4, 64]),
+          [intT(0), intT(0), intT(0), intT(0)],
+      ),
+      (
+          'DUS_kv_example',
+          torch.randn(2, 1280, 4, 64),
+          torch.randn([2, 1024, 4, 64]),
+          [intT(0), intT(0), intT(0), intT(0)],
+      ),
+      (
+          'DUS_3d',
+          torch.randn(2, 256, 4, 64),
+          torch.randn([2, 256, 2, 64]),
+          [intT(0), intT(0), intT(2), intT(0)],
+      ),
+      (
+          'DUS_3d_v2',
+          torch.randn(2, 256, 4, 64),
+          torch.randn([2, 256, 3, 64]),
+          [intT(0), intT(0), intT(1), intT(0)],
+      ),
+      (
+          'DUS_3d_v3',
+          torch.randn(6, 8, 32),
+          torch.randn([6, 3, 32]),
+          [intT(0), intT(5), intT(0)],
+      ),
+      (
+          'DUS_2d',
+          torch.randn(8, 32),
+          torch.randn([8, 12]),
+          [intT(0), intT(20)],
+      ),
+  )
+  def test_opcheck_dynamic_update_slice(self, buffer, update, indices):
+    torch.library.opcheck(dynamic_update_slice, (buffer, update, indices))
+    out = dynamic_update_slice(buffer, update, indices)
+    self.assertTrue(updated_slice_matches(out, update, indices))
+
+  def test_exported_program(self):
+    buffer = torch.randn(1, 1280, 4, 64)
+    update = torch.randn([1, 1024, 4, 64])
+    index = [intT(0), intT(0), intT(0), intT(0)]
+    dm = DUSMod()
+    ep = torch.export.export(dm, (buffer, update, index))
+    dus_in_exported_program = False
+    for node in ep.graph.nodes:
+      if node.op == 'call_function':
+        if node.target.__name__.startswith('dynamic_update_slice'):
+          dus_in_exported_program = True
+          break
+
+    self.assertTrue(dus_in_exported_program)
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_torch/generative/utilities/dynamic_update_slice.py

@@ -0,0 +1,56 @@
+# 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.
+# TODO: Update impl for multiple non-trivial start_indices
+@torch.library.custom_op("ai_edge_torch::dynamic_update_slice", mutates_args=())
+def dynamic_update_slice(
+    in_tensor: torch.Tensor,
+    update: torch.Tensor,
+    start_indices: Sequence[torch.Tensor],
+) -> torch.Tensor:
+  compare_size = torch.tensor(in_tensor.size()) == torch.tensor(update.size())
+  mismatch = torch.nonzero(~compare_size, as_tuple=False)
+  dim = mismatch[0].item() if len(mismatch) > 0 else 0
+  start = start_indices[dim].item()
+  end = start + update.shape[dim]
+  indices = torch.arange(start, end).to(torch.long)
+  return in_tensor.index_copy(dim, indices, update)
+
+
+# Use register_fake to add a ``FakeTensor`` kernel for the operator
+@dynamic_update_slice.register_fake
+def _(in_tensor, update, start_indices):
+  return in_tensor.clone().detach()
+
+
+@lowerings.lower(torch.ops.ai_edge_torch.dynamic_update_slice)
+def _dynamic_update_slice_lower(
+    lctx,
+    in_tensor: ir.Value,
+    update: ir.Value,
+    start_indices: Sequence[ir.Value],
+):
+  return stablehlo.dynamic_update_slice(in_tensor, update, start_indices)

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -185,6 +185,7 @@ def merged_bundle_to_tfl_model(
     converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
     converter._set_original_model_type(conversion_metadata_fb.ModelType.PYTORCH)
     converter._experimental_enable_composite_direct_lowering = True
+    converter._experimental_enable_dynamic_update_slice = True
     converter.model_origin_framework = "PYTORCH"
 
     conversion_utils.set_tfl_converter_quant_flags(converter, quant_config)

ai_edge_torch/odml_torch/jax_bridge/_wrap.py

@@ -24,6 +24,7 @@ from ai_edge_torch.odml_torch.jax_bridge import utils
 import jax
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import func
+from jax._src.lib.mlir.dialects import hlo as stablehlo
 import torch.utils._pytree as pytree
 
 # Jax double (64bit) precision is required to generate StableHLO mlir with
@@ -143,8 +144,39 @@ def wrap(jaxfn: Callable[Any, Any], ir_input_names: list[str] = None):
       ir_inputs = []
 
     results = func.CallOp(cloned_func, ir_inputs).results
+
+    if lctx.node is None:
+      return results[0] if len(results) == 1 else results
+
+    out_avals = lctx.node.meta.get("tensor_meta") or lctx.node.meta.get("val")
+
+    if out_avals is None:
+      return results[0] if len(results) == 1 else results
+
+    def sanitize_result_elty(result, aval):
+      # JAX implementation may not respect aten op's output dtype. For example,
+      # JAX may implement a slightly different dtype upcast rules, leads to
+      # different result's dtype from bridged lowering and torch op output.
+      # Here we add an additional `stablehlo.convert` op when dtype does not
+      # match, to ensure the lowering's result dtype will always be the same
+      # as torch op's output dtype.
+      if aval is None:
+        return result
+
+      target_elty = export_utils.torch_dtype_to_ir_element_type(
+          lctx.ir_context, aval.dtype
+      )
+      if result.type.element_type == target_elty:
+        return result
+      return stablehlo.convert(
+          ir.RankedTensorType.get(result.type.shape, target_elty), result
+      )
+
     if len(results) == 1:
-      return results[0]
-    return results
+      return sanitize_result_elty(results[0], out_avals)
+    return [
+        sanitize_result_elty(result, aval)
+        for result, aval in zip(results, out_avals)
+    ]
 
   return wrapped

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -15,13 +15,17 @@
 import math
 from typing import Optional, Union
 
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch.lowerings import context
+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 numpy as np
 import torch
 
-from .registry import lower
+LoweringContext = context.LoweringContext
+lower = registry.lower
 
 
 # add(Tensor self, Tensor other) -> Tensor
@@ -211,6 +215,31 @@ def _aten_floor(lctx, x: ir.Value, *, out=None) -> ir.Value:
   return stablehlo.floor(x)
 
 
+# Schema:
+#   - aten::cat(Tensor[] tensors, int dim=0) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/main/generated/torch.cat.html
+@lower(torch.ops.aten.cat.default)
+def _aten_cat(lctx: LoweringContext, tensors, dim=0):
+  assert tensors
+  non_empty_tensors = [t for t in tensors if np.prod(t.type.shape) != 0]
+  out_aval = lctx.node.meta.get("tensor_meta") or lctx.node.meta.get("val")
+  if not non_empty_tensors:
+    return utils.splat(
+        0,
+        export_utils.torch_dtype_to_ir_element_type(
+            lctx.ir_context, out_aval.dtype
+        ),
+        out_aval.shape,
+    )
+
+  if dim < 0:
+    dim = dim + len(out_aval.shape)
+  dim = ir.IntegerAttr.get(ir.IntegerType.get_signless(64), dim)
+
+  return stablehlo.concatenate(non_empty_tensors, dim)
+
+
 # Schema:
 #   - aten::slice_scatter(Tensor self, Tensor src, int dim=0, SymInt?
 #       start=None, SymInt? end=None, SymInt step=1) -> Tensor

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -105,7 +105,6 @@ lower_by_torch_xla2(torch.ops.aten.bitwise_not)
 lower_by_torch_xla2(torch.ops.aten.bitwise_or)
 lower_by_torch_xla2(torch.ops.aten.bitwise_xor)
 lower_by_torch_xla2(torch.ops.aten.bmm)
-lower_by_torch_xla2(torch.ops.aten.cat)
 lower_by_torch_xla2(torch.ops.aten.ceil)
 lower_by_torch_xla2(torch.ops.aten.clamp.Tensor)
 lower_by_torch_xla2(torch.ops.aten.clamp.default)
@@ -172,7 +171,6 @@ lower_by_torch_xla2(torch.ops.aten.mm)
 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_backward)
 lower_by_torch_xla2(torch.ops.aten.ne)
 lower_by_torch_xla2(torch.ops.aten.neg)

ai_edge_torch/odml_torch/lowerings/registry.py

@@ -61,6 +61,7 @@ global_registry.decompositions.update(
         torch.ops.aten._adaptive_avg_pool2d,
         torch.ops.aten._adaptive_avg_pool3d,
         torch.ops.aten.grid_sampler_2d,
+        torch.ops.aten.native_group_norm,
         torch.ops.aten.native_dropout,
         torch.ops.aten.reflection_pad1d,
         torch.ops.aten.reflection_pad2d,

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241120.dist-info → ai_edge_torch_nightly-0.3.0.dev20241122.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241120
+Version: 0.3.0.dev20241122
 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
@@ -29,7 +29,7 @@ Requires-Dist: safetensors
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: torch-xla>=2.4.0
-Requires-Dist: tf-nightly>=2.19.0.dev20241001
+Requires-Dist: tf-nightly>=2.19.0.dev20241121
 Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
 

{ai_edge_torch_nightly-0.3.0.dev20241120.dist-info → ai_edge_torch_nightly-0.3.0.dev20241122.dist-info}/RECORD

@@ -3,7 +3,7 @@ ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=52sF7t2CBQE8RcB2Hcmo-f6_BLyCW9NzWZ-wTKM9ho4,706
+ai_edge_torch/version.py,sha256=B4r6opjqsPmDJdLbwvWto6dM-0KbsjszxSL6CXmi8K8,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -27,7 +27,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/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=yXfeWDw9u_rTS3B6kvvFPo5E4XNT3zKTSLFSBSAI9Fc,15502
-ai_edge_torch/_convert/test/test_convert_composites.py,sha256=ELwHxTdTTCJm30aWg_PZXxg9HvDM4Hnf9lT0wwOWT6s,8060
+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
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
@@ -120,7 +120,7 @@ ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-m
 ai_edge_torch/generative/layers/kv_cache.py,sha256=lbm-yJ1jGPtcgWS4C3FmSnB1IlxqDE7g0BLRh3PN4N4,6324
 ai_edge_torch/generative/layers/model_config.py,sha256=viX51T_naJ9sPpPxPoMnSueBPYE2zxWNOD0xn0f-_bM,7510
 ai_edge_torch/generative/layers/normalization.py,sha256=eKAGst9rPuyRFExMcQFJO7R3iHdCtlmjeF_lITjLhwE,6498
-ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=CZqOoibLcHvUgrgaIIWAlmk3XgE2inzx340MN-npLoU,1347
+ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=xxWtlVsGGJkEyXC6PwznubyhJnLPEfSpHOORE_hgxss,2670
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=9jKzOfiBQ66bp1ZnVIAoREIifVNFx4aTlQeYMAx2_pA,29062
@@ -134,6 +134,7 @@ 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_custom_dus.py,sha256=gxG78CcTpXF3iLzDR15Rlz1ey1tNTlSdkp6TeYEijp0,3301
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=W6Bh0gYDzmwb0j9HdD5_D7Z7FPToP2HSyFrmwIXuFqo,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=aZFaheg2sq7rEccch1TZM6W4BSfpJZjrM9Gyp4hVGYs,6351
@@ -142,6 +143,7 @@ ai_edge_torch/generative/test/test_quantize.py,sha256=8geJhKwYBU20m0mdGPD1BUFwQ0
 ai_edge_torch/generative/test/utils.py,sha256=eQ-hjd1eXuHJF3SJK6_CrjgOZVzmG_4VEdH7Z1gH_lA,1897
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/converter.py,sha256=S14STbyxV6A9HKy1BdUo49f2jS6Ij0RL9mVAFUMWYV8,5291
+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=OcHJhEqc3LjI3STli6cyn71m1mdzr7QbzF9fqSNCXrg,5730
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
@@ -158,7 +160,7 @@ ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py,sha256=j8WpeS-mz3Zr4
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
 ai_edge_torch/lowertools/_shim.py,sha256=ilL7x1ebUBj1clg7bagrX4y_nVSHiGrvDrOVfuTeenE,3039
 ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
-ai_edge_torch/lowertools/odml_torch_utils.py,sha256=K5dZ_fFDL3GWKo0IoY4OC_GX5MY-guY-MqteolyV9hg,8098
+ai_edge_torch/lowertools/odml_torch_utils.py,sha256=Smt7p62-lZ_3bBBfnbssAK5GAGxm3U_X7M-1qwsmc68,8161
 ai_edge_torch/lowertools/test_utils.py,sha256=bPgc2iXX16KYtMNvmsRdKfrCY6UJmcfitfCOvHoD7Oc,1930
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=XGZE0vZG9WSQT-6dFmPlU8W89z8rfXPRGjuZeuhXCIw,9205
 ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
@@ -175,16 +177,16 @@ ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=9ag6-WWRG50rPCtIV7OpIokEKu
 ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=1xCXOs3-9UcsOyLFH0uyQwLu7c06iYFTo0NQ7Ckbl2I,1465
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
 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/_wrap.py,sha256=oQo9nxH08NnEDeZaGoCUk1kRtoEOM_f0DUOyd9nfxjg,6673
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=dE_qzh-OnCNjWzqs1-PHs5PNlRF726qMQKM3tkwAzEs,959
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=z_hPJX9n97d6obcsS9OHXpKqbmw6QqACXgnq5ML6Rhs,9014
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=eH9eJqFO-BI9l4WdXfjsItODPRa18SAR_qSvJ6-7gxc,9987
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=v1VdKmL8YLJv3PR9VgyNghO83A25PpTzY2ZUAJqlq3Q,6847
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=1JeX3j7Rt3KE7Z2eYRrhtcYgO3EKnRyZFKAUWXw-bsU,10812
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=4UyNyaR2W-vCOvj-P5lywQ1_RfLIxVE7J_GONI6CQvI,10718
 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=gqx3n1Mx8pnGQz3nkIF1T_8bkRabXLJBvUoJJn5kOUY,2911
+ai_edge_torch/odml_torch/lowerings/registry.py,sha256=itTt8MLbq2LoHTzRidCF2TTbh0TP7L836u99qCjP3FA,2953
 ai_edge_torch/odml_torch/lowerings/utils.py,sha256=NczqpsSd3Fn7yVcPC3qllemiZxxDAZgcW1T5l8-W9fE,5593
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
@@ -194,8 +196,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.dev20241120.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/METADATA,sha256=1Nv_QeerPRw888sOTf4jHx5Ihu-PJD9rL8GOpRHSTa4,1897
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/METADATA,sha256=-YpC-ksRKR8hJ8pZET4Q2F5KbUiRmGOXPhBoEQgIuOA,1897
+ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
+ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/RECORD,,