ai-edge-torch-nightly 0.2.0.dev20240626__py3-none-any.whl → 0.2.0.dev20240701__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -15,6 +15,7 @@
 
 from .convert.converter import convert
 from .convert.converter import signature
+from .convert.to_channel_last_io import to_channel_last_io
 from .model import Model
 
 

ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py

@@ -17,10 +17,22 @@ import functools
 
 import torch
 
-from ai_edge_torch.convert.fx_passes import FxPassBase
-from ai_edge_torch.convert.fx_passes import FxPassResult
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
 from ai_edge_torch.hlfb import mark_pattern
 
+# For torch nightly released after mid June 2024,
+# torch.nn.functional.interpolate no longer gets exported into decomposed graph
+# but single aten op torch.ops.aten.upsample_nearest2d.vec/torch.ops.aten.upsample_bilinear2d.vec.
+# This behavior would our pattern matching based composite builder.
+# It requires the pattern and model graph to get decomposed first for backward compatibility.
+_INTERPOLATE_DECOMPOSITIONS = torch._decomp.get_decompositions(
+    [
+        torch.ops.aten.upsample_bilinear2d.vec,
+        torch.ops.aten.upsample_nearest2d.vec,
+    ]
+)
+
 
 @functools.cache
 def _get_upsample_bilinear2d_pattern():
@@ -30,6 +42,7 @@ def _get_upsample_bilinear2d_pattern():
           x, scale_factor=2, mode="bilinear", align_corners=False
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
+      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -52,6 +65,7 @@ def _get_upsample_bilinear2d_align_corners_pattern():
           x, scale_factor=2, mode="bilinear", align_corners=True
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
+      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -72,6 +86,7 @@ def _get_interpolate_nearest2d_pattern():
       "tfl.resize_nearest_neighbor",
       lambda x: torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest"),
       export_args=(torch.rand(1, 3, 100, 100),),
+      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -86,7 +101,7 @@ def _get_interpolate_nearest2d_pattern():
   return pattern
 
 
-class BuildInterpolateCompositePass(FxPassBase):
+class BuildInterpolateCompositePass(ExportedProgramPassBase):
 
   def __init__(self):
     super().__init__()
@@ -96,10 +111,13 @@ class BuildInterpolateCompositePass(FxPassBase):
         _get_interpolate_nearest2d_pattern(),
     ]
 
-  def call(self, graph_module: torch.fx.GraphModule):
+  def call(self, exported_program: torch.export.ExportedProgram):
+    exported_program = exported_program.run_decompositions(_INTERPOLATE_DECOMPOSITIONS)
+
+    graph_module = exported_program.graph_module
     for pattern in self._patterns:
       graph_module = mark_pattern.mark_pattern(graph_module, pattern)
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return FxPassResult(graph_module, True)
+    return ExportedProgramPassResult(exported_program, True)

ai_edge_torch/convert/test/test_to_channel_last_io.py

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+
+import unittest
+
+import torch
+
+import ai_edge_torch
+
+
+class Identity(torch.nn.Module):
+
+  def forward(self, x):
+    return x
+
+
+class TestToChannelLastIO(unittest.TestCase):
+  """Tests to_channel_last_io API and module wrapper."""
+
+  def test_no_transformations(self):
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity())(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_args(self):
+    x = torch.rand(1, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0])(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_outputs(self):
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity(), outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_args_outputs(self):
+    x = torch.rand(1, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0], outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_args_5d(self):
+    x = torch.rand(1, 10, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0])(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10, 10))
+
+  def test_outputs_5d(self):
+    x = torch.rand(1, 3, 10, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity(), outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 10, 3))
+
+  def test_chained_wrappers(self):
+    x = torch.rand(1, 10, 10, 3)
+
+    m = Identity()
+    m = ai_edge_torch.to_channel_last_io(m, args=[0])
+    m = ai_edge_torch.to_channel_last_io(m, outputs=[0])
+
+    y = m(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_list_args(self):
+    class Add(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    x = (torch.rand(1, 10, 10, 3), torch.rand(1, 10, 10, 3))
+    y = ai_edge_torch.to_channel_last_io(Add(), args=[0, 1])(*x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_list_outputs(self):
+    class TwoIdentity(torch.nn.Module):
+
+      def forward(self, x):
+        return x, x
+
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(TwoIdentity(), outputs=[0])(x)
+    self.assertIsInstance(y, tuple)
+    self.assertEqual(y[0].shape, (1, 10, 10, 3))
+    self.assertEqual(y[1].shape, (1, 3, 10, 10))
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/convert/to_channel_last_io.py

@@ -0,0 +1,85 @@
+# 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.
+# ==============================================================================
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class ChannelLastIOWrapper(nn.Module):
+
+  def __init__(self, wrapped, *, args=None, outputs=None):
+    super().__init__()
+    self.wrapped = wrapped
+    self._args = args or []
+    self._outputs = outputs or []
+
+  def _to_channel_last(self, x):
+    if not torch.is_tensor(x):
+      raise ValueError("Input must be a torch tensor")
+    if x.ndim < 3:
+      raise ValueError("Input must be a tensor with rank >= 3 in layout (N, C, ...)")
+    dims = [0, *range(2, x.ndim), 1]
+    return torch.permute(x, dims)
+
+  def _to_channel_first(self, x):
+    if not torch.is_tensor(x):
+      raise ValueError("Input must be a torch tensor.")
+    if x.ndim < 3:
+      raise ValueError("Input must be a tensor with rank >= 3 in layout (N, ..., C)")
+    dims = [0, x.ndim - 1, *range(1, x.ndim - 1)]
+    return torch.permute(x, dims)
+
+  def forward(self, *args, **kwargs):
+    args = list(args)
+    for i in self._args:
+      args[i] = self._to_channel_first(args[i])
+
+    outputs = self.wrapped(*args, **kwargs)
+
+    if not isinstance(outputs, (list, tuple)):
+      outputs_is_list = False
+      output_list = [outputs]
+    else:
+      outputs_is_list = True
+      output_list = list(outputs)
+
+    for i in self._outputs:
+      output_list[i] = self._to_channel_last(output_list[i])
+
+    if not outputs_is_list:
+      return output_list[0]
+    else:
+      return type(outputs)(output_list)
+
+
+def to_channel_last_io(
+    module: nn.Module,
+    args: Optional[list[int]] = None,
+    outputs: Optional[list[int]] = None,
+):
+  """Wraps the module with channel first to channel last layout transformations.
+
+  Args:
+    args (list[int]): Transform args with indices in the list from channel first
+      (N, C, ...) to channel last (N, ..., C).
+    outputs (list[int]): Transform outputs with indices in the list from channel
+      first (N, C, ...) to channel last (N, ..., C).
+  Returns:
+    The wrapped nn.Module with additional layout transposes after inputs and/or before
+    outputs.
+  """
+  return ChannelLastIOWrapper(module, args=args, outputs=outputs)

ai_edge_torch/generative/examples/gemma/convert_to_tflite.py

@@ -45,9 +45,9 @@ def convert_gemma_to_tflite(
       checkpoint_path, kv_cache_max_len=kv_cache_max_len
   )
   # Tensors used to trace the model graph during conversion.
-  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.int)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
-  decode_token = torch.tensor([[0]], dtype=torch.long)
+  decode_token = torch.tensor([[0]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
 
   quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else None

ai_edge_torch/generative/examples/gemma/gemma.py

@@ -163,7 +163,7 @@ def define_and_run_2b() -> None:
   checkpoint_path = os.path.join(Path.home(), "Downloads/llm_data/gemma-2b")
   model = build_2b_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx
   input_pos = torch.arange(0, kv_cache_max_len)
   print("running an inference")

ai_edge_torch/generative/examples/phi2/convert_to_tflite.py

@@ -43,9 +43,9 @@ def convert_phi2_to_tflite(
   """
   pytorch_model = phi2.build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
   # Tensors used to trace the model graph during conversion.
-  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.int)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
-  decode_token = torch.tensor([[0]], dtype=torch.long)
+  decode_token = torch.tensor([[0]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
 
   quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else None

ai_edge_torch/generative/examples/phi2/phi2.py

@@ -153,7 +153,7 @@ def define_and_run() -> None:
   checkpoint_path = os.path.join(Path.home(), "Downloads/llm_data/phi2")
   model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx
   input_pos = torch.arange(0, kv_cache_max_len)
   print("running an inference")

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -60,8 +60,8 @@ class CLIP(nn.Module):
     )
 
   @torch.inference_mode
-  def forward(self, tokens: torch.LongTensor) -> torch.FloatTensor:
-    tokens = tokens.type(torch.long)
+  def forward(self, tokens: torch.IntTensor) -> torch.FloatTensor:
+    tokens = tokens.type(torch.int)
 
     state = self.tok_embedding(tokens) + self.tok_embedding_position
     for layer in self.transformer_blocks:

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -61,7 +61,7 @@ def convert_stable_diffusion_to_tflite(
   n_tokens = 77
   timestamp = 0
   len_prompt = 1
-  prompt_tokens = torch.full((1, n_tokens), 0, dtype=torch.long)
+  prompt_tokens = torch.full((1, n_tokens), 0, dtype=torch.int)
   input_image = torch.full((1, 3, image_height, image_width), 0, dtype=torch.float32)
   noise = torch.full(
       (len_prompt, 4, image_height // 8, image_width // 8), 0, dtype=torch.float32

ai_edge_torch/generative/examples/t5/convert_to_tflite.py

@@ -30,23 +30,23 @@ def convert_t5_to_tflite_singlesig(checkpoint_path: str):
 
   # encoder
   seq_len = 512
-  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
   prompt_e_token = [1, 2, 3, 4, 5, 6]
   prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(
-      prompt_e_token, dtype=torch.long
+      prompt_e_token, dtype=torch.int
   )
   prefill_e_input_pos = torch.arange(0, seq_len)
-  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
   prompt_d_token = [1, 2, 3, 4, 5, 6]
   prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(
-      prompt_d_token, dtype=torch.long
+      prompt_d_token, dtype=torch.int
   )
   prefill_d_input_pos = torch.arange(0, seq_len)
 
   # decoder
-  decode_token = torch.tensor([[1]], dtype=torch.long)
+  decode_token = torch.tensor([[1]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
-  decode_d_token = torch.tensor([[1]], dtype=torch.long)
+  decode_d_token = torch.tensor([[1]], dtype=torch.int)
   decode_d_input_pos = torch.tensor([0], dtype=torch.int64)
 
   # Pad mask for self attention only on "real" tokens.
@@ -78,23 +78,23 @@ def convert_t5_to_tflite_multisig(checkpoint_path: str):
 
   # encoder
   seq_len = 512
-  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
   prompt_e_token = [1, 2, 3, 4, 5, 6]
   prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(
-      prompt_e_token, dtype=torch.long
+      prompt_e_token, dtype=torch.int
   )
   prefill_e_input_pos = torch.arange(0, seq_len)
-  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
   prompt_d_token = [1, 2, 3, 4, 5, 6]
   prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(
-      prompt_d_token, dtype=torch.long
+      prompt_d_token, dtype=torch.int
   )
   prefill_d_input_pos = torch.arange(0, seq_len)
 
   # decoder
-  decode_token = torch.tensor([[1]], dtype=torch.long)
+  decode_token = torch.tensor([[1]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
-  decode_d_token = torch.tensor([[1]], dtype=torch.long)
+  decode_d_token = torch.tensor([[1]], dtype=torch.int)
   decode_d_input_pos = torch.tensor([0], dtype=torch.int64)
 
   # Pad mask for self attention only on "real" tokens.

ai_edge_torch/generative/examples/t5/t5.py

@@ -562,7 +562,7 @@ def define_and_run_t5(checkpoint_path: str) -> None:
   model = build_t5_model(checkpoint_path)
 
   idx = get_sample_encoder_input_ids()
-  tokens = torch.full((1, 512), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, 512), 0, dtype=torch.int, device="cpu")
   tokens[0, :77] = idx
   input_pos = torch.arange(0, 512)
 
@@ -586,7 +586,7 @@ def define_and_run_t5_split(checkpoint_path: str) -> None:
   t5_decoder_model = build_t5_decoder_model(config, embedding_layer, checkpoint_path)
   idx = get_sample_encoder_input_ids()
 
-  tokens = torch.full((1, 512), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, 512), 0, dtype=torch.int, device="cpu")
   tokens[0, :77] = idx
   input_pos = torch.arange(0, 512)
 

ai_edge_torch/generative/examples/test_models/toy_model.py

@@ -93,7 +93,7 @@ def define_and_run() -> None:
   )
 
   model = ToySingleLayerModel(config)
-  idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN), 0)
+  idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN, dtype=torch.int), 0)
   input_pos = torch.arange(0, KV_CACHE_MAX_LEN)
   print('running an inference')
   print(

ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py

@@ -115,13 +115,13 @@ def get_model_config() -> cfg.ModelConfig:
 
 
 def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
-  idx = torch.unsqueeze(torch.arange(0, 100), 0)
+  idx = torch.unsqueeze(torch.arange(0, 100, dtype=torch.int), 0)
   input_pos = torch.arange(0, 100)
   return idx, input_pos
 
 
 def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
-  idx = torch.tensor([[1]], dtype=torch.long)
+  idx = torch.tensor([[1]], dtype=torch.int)
   input_pos = torch.tensor([10])
   return idx, input_pos
 

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -103,13 +103,13 @@ def get_model_config() -> cfg.ModelConfig:
 
 
 def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
-  idx = torch.unsqueeze(torch.arange(0, 100), 0)
+  idx = torch.unsqueeze(torch.arange(0, 100, dtype=torch.int), 0)
   input_pos = torch.arange(0, 100)
   return idx, input_pos
 
 
 def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
-  idx = torch.tensor([[1]], dtype=torch.long)
+  idx = torch.tensor([[1]], dtype=torch.int)
   input_pos = torch.tensor([10], dtype=torch.int64)
   return idx, input_pos
 

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -45,9 +45,9 @@ def convert_tiny_llama_to_tflite(
       checkpoint_path, kv_cache_max_len=kv_cache_max_len
   )
   # Tensors used to trace the model graph during conversion.
-  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.int)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
-  decode_token = torch.tensor([[0]], dtype=torch.long)
+  decode_token = torch.tensor([[0]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
 
   quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else None

ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py

@@ -153,7 +153,7 @@ def define_and_run() -> None:
   checkpoint_path = os.path.join(Path.home(), "Downloads/llm_data/tiny_llama")
   model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx
   input_pos = torch.arange(0, kv_cache_max_len)
   print("running an inference")

ai_edge_torch/generative/quantize/example.py

@@ -26,7 +26,7 @@ def main():
   config = gemma.get_fake_model_config_2b_for_test()
   model = gemma.Gemma(config)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx
   input_pos = torch.arange(0, 10)
 

ai_edge_torch/generative/test/test_model_conversion.py

@@ -35,7 +35,7 @@ class TestModelConversion(unittest.TestCase):
   def test_toy_model_with_kv_cache(self):
     config = toy_model_with_kv_cache.get_model_config()
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int64
     )
 
@@ -59,7 +59,7 @@ class TestModelConversion(unittest.TestCase):
     config = toy_model_with_kv_cache.get_model_config()
     config.batch_size = 2
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
-    idx, input_pos = torch.tensor([[1], [2]], dtype=torch.long), torch.tensor(
+    idx, input_pos = torch.tensor([[1], [2]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int64
     )
 
@@ -83,7 +83,7 @@ class TestModelConversion(unittest.TestCase):
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int64
     )
 
@@ -109,7 +109,7 @@ class TestModelConversion(unittest.TestCase):
     pytorch_model = tiny_llama.TinyLLamma(config)
 
     idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+    tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
     tokens[0, :4] = idx
     input_pos = torch.arange(0, 10)
 
@@ -135,13 +135,13 @@ class TestModelConversion(unittest.TestCase):
 
     # prefill
     seq_len = 10
-    prefill_tokens = torch.full((1, seq_len), 0, dtype=torch.long, device="cpu")
+    prefill_tokens = torch.full((1, seq_len), 0, dtype=torch.int, device="cpu")
     prompt_token = torch.from_numpy(np.array([1, 2, 3, 4]))
     prefill_tokens[0, : len(prompt_token)] = prompt_token
     prefill_input_pos = torch.arange(0, seq_len)
 
     # decode
-    decode_token = torch.tensor([[1]], dtype=torch.long)
+    decode_token = torch.tensor([[1]], dtype=torch.int)
     decode_input_pos = torch.tensor([5], dtype=torch.int64)
 
     edge_model = (
@@ -183,7 +183,7 @@ class TestModelConversion(unittest.TestCase):
     model = gemma.Gemma(config)
 
     idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+    tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
     tokens[0, :4] = idx
     input_pos = torch.arange(0, 10)
 
@@ -210,7 +210,7 @@ class TestModelConversion(unittest.TestCase):
     pytorch_model = phi2.Phi2(config)
 
     idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+    tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
     tokens[0, :4] = idx
     input_pos = torch.arange(0, 10)
 

ai_edge_torch/generative/test/test_quantize.py

@@ -119,7 +119,7 @@ class TestQuantizeConvert(unittest.TestCase):
     self.skipTest("b/346896669")
     config = toy_model_with_kv_cache.get_model_config()
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int64
     )
 
@@ -137,7 +137,7 @@ class TestQuantizeConvert(unittest.TestCase):
     self.skipTest("b/338288901")
     config = toy_model_with_kv_cache.get_model_config()
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int64
     )
 

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -100,13 +100,18 @@ class ScalarAttrLocation:
 
 
 def _find_scalar_attr(
-    pattern_module: torch.nn.Module, export_args: tuple[Any], tracker: ScalarAttrTracker
+    pattern_module: torch.nn.Module,
+    export_args: tuple[Any],
+    tracker: ScalarAttrTracker,
+    decomp_table=None,
 ) -> ScalarAttrLocation:
   scalar_loc_intersections = None
   for source, target in tracker._source_targets:
     track_args = list(export_args)
     track_args[tracker.pattern_arg_pos] = source
     ep = torch.export.export(pattern_module, tuple(track_args))
+    if decomp_table is not None:
+      ep = ep.run_decompositions(decomp_table)
 
     scalar_locs = set()
     nodes = ep.graph_module.graph.nodes
@@ -145,6 +150,7 @@ class Pattern:
           ["Pattern", GraphModule, InternalMatch], Optional[dict[str, Any]]
       ] = None,
       scalar_attr_trackers: list[ScalarAttrTracker] = None,
+      decomp_table: Optional[dict[torch._ops.OperatorBase, Callable]] = None,
   ):
     """The PyTorch computation pattern to match against a model.
 
@@ -165,6 +171,8 @@ class Pattern:
         for scalar args in `export_args`, which are used to track
         the attr occurrence(s) and retrieve their values from the
         matched subgraph.
+      decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]):
+        The decomposition table to be run on the pattern's exported program.
     """
     if not isinstance(module, torch.nn.Module):
 
@@ -180,23 +188,28 @@ class Pattern:
       module = PatternModule(module).eval()
 
     self.name = name
-    self.exported_program = torch.export.export(module, export_args)
-    self.graph_module = self.exported_program.graph_module
     self.attr_builder = attr_builder
     self._scalar_attr_trackers = scalar_attr_trackers if scalar_attr_trackers else []
 
-    # Sanitize graph_module for more precise pattern matching.
-    # The graph_module to match against this pattern should apply equivalent
-    # sanitization.
-    self.graph_module = passes.remove_clone_ops(self.graph_module)
-    self.graph_module = passes.remove_dangling_args(self.graph_module)
+    exported_program = torch.export.export(module, export_args)
+    if decomp_table is not None:
+      exported_program = exported_program.run_decompositions(decomp_table)
+
+    self.exported_program = exported_program
+    self.graph_module = self.exported_program.graph_module
 
     self._scalar_attr_locations = []
     for tracker in self._scalar_attr_trackers:
       self._scalar_attr_locations.append(
-          _find_scalar_attr(module, export_args, tracker)
+          _find_scalar_attr(module, export_args, tracker, decomp_table=decomp_table)
       )
 
+    # Sanitize graph_module for more precise pattern matching.
+    # The graph_module to match against this pattern should apply equivalent
+    # sanitization.
+    self.graph_module = passes.remove_clone_ops(self.graph_module)
+    self.graph_module = passes.remove_dangling_args(self.graph_module)
+
     # Builds list of ordered input and output nodes.
     self.graph_nodes_map = {}
     for node in self.graph_module.graph.nodes:

{ai_edge_torch_nightly-0.2.0.dev20240626.dist-info → ai_edge_torch_nightly-0.2.0.dev20240701.dist-info}/METADATA

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

@@ -1,13 +1,14 @@
-ai_edge_torch/__init__.py,sha256=FPMmuFU3pyMREtjB_san1fy_0PFtAsgA0VZfOYvDrb4,1008
+ai_edge_torch/__init__.py,sha256=CNDboRP4zQBpz2hznNCQWcQCARvNXUm3DMa1Dw_XXFg,1067
 ai_edge_torch/model.py,sha256=kmcgELjsYl8YzF8nUF6P7q4i8MWS-pLGpfsy-yTUXmE,4243
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/convert/conversion.py,sha256=8K8jQuaCjlUWoj7jiimxp_zpN6mYThLOcQ858UDcYnE,4159
 ai_edge_torch/convert/conversion_utils.py,sha256=9BqCL38DErv1vEVGtT3BIJVhdwZjw2EQ-_m5UpvVVYE,11294
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
+ai_edge_torch/convert/to_channel_last_io.py,sha256=zo5tY3yDhY_EPCkrL1XSXs2uRFS8B4_qu08dSjNsUGk,2778
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=ijVyDclPnd6a0DWWUJkwR4igj6f82S-cE1-83QGPvgw,1652
 ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=2yqUwJJ2R233_X9FNMOP9oYRTTzH34TR_BIUj-wfnKw,7080
-ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=jB27GlDC8x36nn35aiq97uKERiq4KXSUZ7tv7yc0Gl4,3223
+ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=6m_vcycd9f3OQgQLx2hhQjsKfOqdxE5EkjzqrxqyAQM,4168
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=UX6dJsxCqSkftXXvNBV-i7Bjk6H7qTyqzUnE640Itfg,1673
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=VA9bekxPVhLk4MYlIRXnOzrSnbCtUmGj7OQ_fJcKQtc,795
@@ -24,6 +25,7 @@ ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
 ai_edge_torch/convert/test/test_convert_composites.py,sha256=_Ojc-H6GOS5s8ek3_8eRBL_AiCs-k3srziPJ2R4Ulrg,7255
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
+ai_edge_torch/convert/test/test_to_channel_last_io.py,sha256=I8c4ZG3v1vo0yxQYzLK_BTId4AOL9vadHGDtfCUZ4UI,2930
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
 ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
 ai_edge_torch/debug/utils.py,sha256=hjVmQVVl1dKxEF0D6KB4a3ouQ3wBkTsebOX2YsUObZM,1430
@@ -34,15 +36,15 @@ ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 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/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=dZv3r24uHsTMokEdnl3nf7LpmV0q7FLnVtCuHn5AuUs,2538
-ai_edge_torch/generative/examples/gemma/gemma.py,sha256=1lZfXGHmbII4rFu0U2B9NzlJCRhphxtmQtkCHQ39_uw,5935
+ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=yl36sMjVqDlf9I41DF9C5wx6ztMxYB5xukD1NltUS04,2536
+ai_edge_torch/generative/examples/gemma/gemma.py,sha256=caBQrJTK4tXFeGM-i2cNXd2Tb8GXi463MwFHc7N65WE,5934
 ai_edge_torch/generative/examples/phi2/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=6nOuwx9q3AUlYcjXRRXSr_3M2JKqdJ-vUf-uE3VFYHE,2512
-ai_edge_torch/generative/examples/phi2/phi2.py,sha256=PMhKC6JCAMYSj2F3UmWHWK4rTcXD-B6PuehaoDccRqk,5562
+ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=jIFnJY9BtAUJtVkFnms7byZp-jshhQIx59DmK0OjJ8M,2510
+ai_edge_torch/generative/examples/phi2/phi2.py,sha256=tYo-WWOyw6LBF5wBVL_CsuTxKhur0SroDe2sYkTPdvI,5561
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=Lo4Dq7a3Kg-lyH56iqGtqCo5UaClQHRCTDdNagXGTo8,3535
-ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=qU1wVEcn_biwCuDguZljhlLGzpLIqgqC31Dh_lXquQc,3720
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=wVEjsKd5JCIiYf5GF19rOXs2NHscZh0D69mxaS4f0Sk,4182
+ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=cIOqRZ76Pb8ywuCa3LUQnKnBVmvcaAPqvA5bdHfgaWw,3718
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=AwGVIY_tNlea1c4Rz3lEAeAqBvb-GxGzleX2dle98DE,4181
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=RgxedILk7iNMb0mhE4VkCs6d7BnFzYhR3vspUkC0-1o,11425
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=sRevfsmCun7zbceJbOstLKNUsLwzQDsGm7Mi2JmlREg,26021
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=mgbxkeFDMkNIGmnbcFTIFPu8EWKokghiviYIOB2lE3Q,3437
@@ -55,16 +57,16 @@ ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler_ancestral.py
 ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py,sha256=iPYX9ZSaxwSak2KI44j6TEr_g4pdxS3xpka4u0trjbo,2788
 ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py,sha256=5iRfU5MO6GR6K3WrdddIU_9U7ZZGEEb7zGKVY1WFl-8,1340
 ai_edge_torch/generative/examples/t5/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=bWtwtUacvJOEDUpuYvLTgkP7oTkXKJA-Tf4FPxlD1Cw,4536
-ai_edge_torch/generative/examples/t5/t5.py,sha256=L6YrVzUEzP-Imb8W28LdukFGrx1aWSzz1kyYK_9RFZM,21087
+ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=vqYip3JnjFMN5s0VnIzwhlBJl3up75WMb_VtRTfMOK0,4524
+ai_edge_torch/generative/examples/t5/t5.py,sha256=TyBDb50NbtKaHyGLGOMJ8dR2_GPF43oe8WhRfb1SMZ4,21085
 ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=rkMwi-NJGBXHm5S57Rsj1LbcoVdyRkS7GmIBuU6F_2E,8274
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=CUXsQ_IU96NaCg9jyfeKI0Zz2iWDkJUsPJyPR1Pgz7I,3813
-ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py,sha256=zwCmCnhr-vhBwHqv9i7xMasdBGVNqAGxZvWsncsJn58,5543
-ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=lfYUiem_Pbn3vGgPx84BeI8n7rN3-1fImwCLm8Eo2U8,4853
+ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=UNj7IgyCLOlaOu2xJKHciHqcd-_NXLNwMphCWwuaCes,3830
+ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py,sha256=u1M38wuev-YG-exMG0_HzAw3yaikPRhEZCqGaHO7Cw0,5559
+ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=xy7kT1f8V0jqAfglggoeKVqR5KviKdfQun0WBYJ0jS8,4869
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=E4I5OlC4zyl5cxiiu7uTED-zcwYRu210lP1zuT3xLBE,2566
-ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=IFRLPG9wz_aLl_zV_6CETCjSM03ukA6bZqqyDLVACuw,5651
+ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=rAsrP0CnhcYYMCcEJAGoFMRMur0LPoAcYpjFlBiZd2s,2564
+ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=__9QjZ-7Xi8MX5M7iRvw_jpic0-Fau2IO0QatleUzL0,5650
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=aXvYiaHDvETIrh0Q9DDZA_ZBiazGk80DT6nt7lLtC1o,1172
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=IehLwFNwa0C9fnk1pmNmyfuAwwWbuwdyKy46BSqNVdI,1948
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -82,7 +84,7 @@ ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=H45wsXA6iJi_Mjd66NiQrh7
 ai_edge_torch/generative/layers/unet/builder.py,sha256=NmJiZ2-e1wbv9jnvI3VCyUJlONV5ZAOz-RTc7ipAZ5U,1872
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=FrIO-CR8aRIV2i8aFqom_4S7WCEDLMyYwo6U0oFyn7A,9097
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/quantize/example.py,sha256=t-YwyKSPAG-OZC1DfH-0vfie2RHHpTSQjxUY-tmhu5g,1543
+ai_edge_torch/generative/quantize/example.py,sha256=zgBgMyZ8RlSIjRhbaeodLyt6sj_dYuM7oI6Zyx7xQIg,1542
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=Y8zahKw7b_h7ajPaJZVef4jG-MoqImRCpVSbFtV_i24,5139
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=-vd6Qp0BdXJVKg4f0_hhwbKOi3QPIAPVqyXnJ-ZnISQ,1915
@@ -92,8 +94,8 @@ ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py,sha256=47DE
 ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py,sha256=qUB4f2DoB14dLkNPWf6TZodpT81mfAJeWM-lCAmkuHY,5735
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/loader_test.py,sha256=N88CbrLW7Q2x1EyurwdXQ6YjsA-ySQcPxpZH3QOGp-M,3317
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=LsPTrLC1I4JW2GowTS3V9Eu257vLHr2Yj5f_qaFUX84,7589
-ai_edge_torch/generative/test/test_quantize.py,sha256=IjCbCPWzIgXk3s7y7SJsg2usIxhOqs3PuhFvEYR4Sdw,5388
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=tXES8gePl8BptETyUgpIznbIK1SvGEd8mKz1bT_a8Mw,7581
+ai_edge_torch/generative/test/test_quantize.py,sha256=Emp-8oLHyGddELSCkncuSQt8ZJIhZ2-y0-ghR92s10g,5386
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/loader.py,sha256=Hs92478j1g4jQGvbdP1aWvOy907HjwqQZE-NFy6HELo,11326
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=7ChqrnthD7I-Be6vkRvYTRhbGQ3tqMbikLpjY5HpSzE,30890
@@ -101,7 +103,7 @@ ai_edge_torch/generative/utilities/t5_loader.py,sha256=h1FQzt4x8wiQMX4NzYNVIaJGL
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=YV2YKBkh7y7j7sd7EA81vf_1hUKUvTRiy1pfqZustXc,1539
-ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=H4047w-xwx27rYPKNqmeOSQ9M1Adkpd7drp81YdV7Hw,9206
+ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=BEMKgkX8IrsX70h2CxwA_tpsBm6qWWe2bOeOufMYNkA,9722
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=qYR3PRGS9W3OG-qX7UFqL69VxXuUSfyDBUJtCXtXcOE,4262
 ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py,sha256=aUAPKnH4_Jxpp3mLlD5rzdT1g_VBm7OrwwLJ9DeJlzQ,8190
@@ -112,8 +114,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 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=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/METADATA,sha256=taaGTe-WdFG7HWUB43xnjYmG_iFcRwaIaZCK91-QP7M,1748
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240701.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240701.dist-info/METADATA,sha256=QItv3j92_LbO-UPcMqZHRhwTeSTnVQiy-cam2ciLaGM,1748
+ai_edge_torch_nightly-0.2.0.dev20240701.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240701.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240701.dist-info/RECORD,,