ai-edge-torch-nightly 0.2.0.dev20240604__py3-none-any.whl → 0.2.0.dev20240606__py3-none-any.whl

ai_edge_torch/generative/test/test_model_conversion.py

@@ -33,7 +33,6 @@ class TestModelConversion(unittest.TestCase):
   """Unit tests that check for model conversion and correctness."""
 
   def test_toy_model_with_kv_cache(self):
-    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(
@@ -42,19 +41,21 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
 
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            pytorch_model,
-            (idx, input_pos),
-            num_valid_inputs=1,
-            atol=1e-5,
-            rtol=1e-5,
-        )
-    )
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (idx, input_pos),
+              num_valid_inputs=1,
+              atol=1e-5,
+              rtol=1e-5,
+          )
+      )
 
   def test_toy_model_with_kv_cache_with_hlfb(self):
-    self.skipTest("b/338288901")
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
@@ -64,16 +65,19 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
 
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            pytorch_model,
-            (idx, input_pos),
-            num_valid_inputs=1,
-            atol=1e-5,
-            rtol=1e-5,
-        )
-    )
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (idx, input_pos),
+              num_valid_inputs=1,
+              atol=1e-5,
+              rtol=1e-5,
+          )
+      )
 
   def test_tiny_llama(self):
     self.skipTest("b/338288901")
@@ -87,19 +91,21 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
 
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            pytorch_model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
-            atol=1e-5,
-            rtol=1e-5,
-        )
-    )
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (tokens, input_pos),
+              num_valid_inputs=1,
+              atol=1e-5,
+              rtol=1e-5,
+          )
+      )
 
   def test_tiny_llama_multisig(self):
-    self.skipTest("b/338288901")
     config = tiny_llama.get_fake_model_config_for_test()
     pytorch_model = tiny_llama.TinyLLamma(config)
 
@@ -122,32 +128,30 @@ class TestModelConversion(unittest.TestCase):
         .convert()
     )
 
-    # For the pytorch model, the KV cache is a persistent state internal to the model, and it
-    # will be shared for prefill and decode. However, for tflite, currently we can't share
-    # kv-cache between the two signatures. prefill will change the content in kv-cache,
-    # but it won't be readable by the decode tflite model. This means the output of running `decode` after
-    # running `prefill` in pytorch will be different from the output of running `decode` after `prefill` via ai_edge_torch.
-    copied_model = copy.deepcopy(pytorch_model)
-
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            pytorch_model,
-            (prefill_tokens, prefill_input_pos),
-            signature_name="prefill",
-            num_valid_inputs=1,
-        )
-    )
-
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            copied_model,
-            (decode_token, decode_input_pos),
-            signature_name="decode",
-            num_valid_inputs=1,
-        )
-    )
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      copied_model = copy.deepcopy(pytorch_model)
+
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (prefill_tokens, prefill_input_pos),
+              signature_name="prefill",
+              num_valid_inputs=1,
+          )
+      )
+
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              copied_model,
+              (decode_token, decode_input_pos),
+              signature_name="decode",
+              num_valid_inputs=1,
+          )
+      )
 
   def test_gemma(self):
     self.skipTest("b/338288901")
@@ -161,17 +165,20 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
 
-    # TODO(talumbau, haoliang): debug numerical diff.
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
-            atol=1e-2,
-            rtol=1e-5,
-        )
-    )
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      # TODO(talumbau, haoliang): debug numerical diff.
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              model,
+              (tokens, input_pos),
+              num_valid_inputs=1,
+              atol=1e-2,
+              rtol=1e-5,
+          )
+      )
 
   def test_phi2(self):
     self.skipTest("b/338288901")
@@ -185,16 +192,19 @@ class TestModelConversion(unittest.TestCase):
 
     edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
 
-    self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            edge_model,
-            pytorch_model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
-            atol=1e-5,
-            rtol=1e-5,
-        )
-    )
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (tokens, input_pos),
+              num_valid_inputs=1,
+              atol=1e-5,
+              rtol=1e-5,
+          )
+      )
 
 
 if __name__ == "__main__":

ai_edge_torch/generative/utilities/loader.py

@@ -69,10 +69,16 @@ def load_pytorch_statedict(full_path: str):
   Raises:
     ValueError: If no tensors are loaded from the provided directory or file.
   """
-  pattern = os.path.join(full_path, "*.bin") if os.path.isdir(full_path) else full_path
   files = []
-  for file in glob.glob(pattern):
-    files.append(file)
+  patterns = []
+  if os.path.isdir(full_path):
+    patterns.append(os.path.join(full_path, "*.bin"))
+    patterns.append(os.path.join(full_path, "*.pt"))
+  else:
+    patterns.append(full_path)
+  for pattern in patterns:
+    for file in glob.glob(pattern):
+      files.append(file)
 
   tensors = {}
   for file in files:
@@ -93,18 +99,20 @@ class ModelLoader:
 
   @dataclass
   class TensorNames:
-    attn_query_proj: str
-    attn_key_proj: str
-    attn_value_proj: str
-    attn_output_proj: str
-
-    ff_up_proj: str
-    ff_down_proj: str
+    attn_query_proj: str = None
+    attn_key_proj: str = None
+    attn_value_proj: str = None
+    attn_fused_qkv_proj: str = None
+    attn_output_proj: str = None
+
+    ff_up_proj: str = None
+    ff_down_proj: str = None
     ff_gate_proj: str = None
 
     pre_attn_norm: str = None
     pre_ff_norm: str = None
     embedding: str = None
+    embedding_position: str = None
     final_norm: str = None
     lm_head: str = None
 
@@ -129,6 +137,10 @@ class ModelLoader:
         strict (bool, optional): Whether the converted keys are strictly
           matched. Defaults to True.
 
+    Returns:
+        missing_keys (List[str]): a list of str containing the missing keys
+        unexpected_keys (List[str]): a list of str containing the unexpected keys
+
     Raises:
         ValueError: If conversion results in unmapped tensors and strict mode is
           enabled.
@@ -139,6 +151,10 @@ class ModelLoader:
       converted_state["tok_embedding.weight"] = state.pop(
           f"{self._names.embedding}.weight"
       )
+      if self._names.embedding_position is not None:
+        converted_state["tok_embedding_position"] = state.pop(
+            f"{self._names.embedding_position}"
+        )
     if self._names.lm_head is not None:
       converted_state["lm_head.weight"] = state.pop(f"{self._names.lm_head}.weight")
       if model.config.lm_head_use_bias:
@@ -158,7 +174,7 @@ class ModelLoader:
       raise ValueError(
           f"Failed to map all tensor. Remaing tensor are: {list(state.keys())}"
       )
-    model.load_state_dict(converted_state, strict=strict)
+    return model.load_state_dict(converted_state, strict=strict)
 
   def _get_loader(self) -> Callable[[str], Dict[str, torch.Tensor]]:
     """A best effort method for finding appropriate state loader.
@@ -172,13 +188,15 @@ class ModelLoader:
     if os.path.isdir(self._file_name):
       if glob.glob(os.path.join(self._file_name, "*.safetensors")):
         return load_safetensors
-      if glob.glob(os.path.join(self._file_name, "*.bin")):
+      if glob.glob(os.path.join(self._file_name, "*.bin")) or glob.glob(
+          os.path.join(self._file_name, "*.pt")
+      ):
         return load_pytorch_statedict
 
     if self._file_name.endswith(".safetensors"):
       return load_safetensors
 
-    if self._file_name.endswith(".bin"):
+    if self._file_name.endswith(".bin") or self._file_name.endswith(".pt"):
       return load_pytorch_statedict
 
     raise ValueError(f"File format not supported.")
@@ -225,22 +243,33 @@ class ModelLoader:
       converted_state: Dict[str, torch.Tensor],
   ):
     prefix = f"transformer_blocks.{idx}"
-    q_name = self._names.attn_query_proj.format(idx)
-    k_name = self._names.attn_key_proj.format(idx)
-    v_name = self._names.attn_value_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = self._fuse_qkv(
-        config,
-        state.pop(f"{q_name}.weight"),
-        state.pop(f"{k_name}.weight"),
-        state.pop(f"{v_name}.weight"),
-    )
-    if config.attn_config.qkv_use_bias:
-      converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = self._fuse_qkv(
+    if self._names.attn_fused_qkv_proj:
+      fused_qkv_name = self._names.attn_fused_qkv_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = state.pop(
+          f"{fused_qkv_name}.weight"
+      )
+    else:
+      q_name = self._names.attn_query_proj.format(idx)
+      k_name = self._names.attn_key_proj.format(idx)
+      v_name = self._names.attn_value_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = self._fuse_qkv(
           config,
-          state.pop(f"{q_name}.bias"),
-          state.pop(f"{k_name}.bias"),
-          state.pop(f"{v_name}.bias"),
+          state.pop(f"{q_name}.weight"),
+          state.pop(f"{k_name}.weight"),
+          state.pop(f"{v_name}.weight"),
       )
+    if config.attn_config.qkv_use_bias:
+      if self._names.attn_fused_qkv_proj:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = state.pop(
+            f"{fused_qkv_name}.bias"
+        )
+      else:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = self._fuse_qkv(
+            config,
+            state.pop(f"{q_name}.bias"),
+            state.pop(f"{k_name}.bias"),
+            state.pop(f"{v_name}.bias"),
+        )
 
     o_name = self._names.attn_output_proj.format(idx)
     converted_state[f"{prefix}.atten_func.output_projection.weight"] = state.pop(

{ai_edge_torch_nightly-0.2.0.dev20240604.dist-info → ai_edge_torch_nightly-0.2.0.dev20240606.dist-info}/METADATA

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

@@ -6,7 +6,7 @@ ai_edge_torch/convert/conversion_utils.py,sha256=NpVm3Ms81_cIW5IYgGsr0BVganJJgBK
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=Ll2nNwufjcV5nSruQPXiloq7F1E7pWJ2T5clXmy1lk8,2825
 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=wHVWNNMu5h_ya6GnnJn0cNif9xmdSqr8Vm-R7lllxZM,6213
+ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=2yqUwJJ2R233_X9FNMOP9oYRTTzH34TR_BIUj-wfnKw,7080
 ai_edge_torch/convert/fx_passes/build_upsample_bilinear2d_composite_pass.py,sha256=76XYoIlFDgrzp5QemoaEalPFcEbfszkEH_PLvO1ASCk,2607
 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
@@ -22,7 +22,7 @@ ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partition
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=FlNKt2EhIKnlVEeUWTiv5sz446YKU6Yy1H0Gd6VRgkU,6432
 ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
-ai_edge_torch/convert/test/test_convert_composites.py,sha256=SrVn_cEMtQhYYCMOUKK0K7M57MQNQX-lOUwieln0HGA,6616
+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/debug/__init__.py,sha256=TKvmnjVk3asvYcVh6C-LPr6srgAF_nppSAupWEXqwPY,707
 ai_edge_torch/debug/culprit.py,sha256=vklaxBUfINdo44OsH7csILK70N41gEThCGchGEfbTZw,12789
@@ -40,10 +40,10 @@ ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=6nOuwx9q3AUlY
 ai_edge_torch/generative/examples/phi2/phi2.py,sha256=VvigzPQ_LJHeADTsMliwFwPe2BcnOhFgKDqr_WZ2JQ8,5540
 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=KR1Ci4rlJeeGfsFRliCxUve9K7RTJLZfTRMgFtfQ4MU,2434
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=6REAYy1Bv-Iv5zcmA_m_W6fH6jt5a3IS6Vge18jS_Wo,3633
+ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=yUCJemEh4n8ez-yLgVU0HZAki-PZ9nY04DFjgpx9PUc,3698
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=MI73RjOeD4Kh7AL0j5_QXiZq-rl_qCdibSE6eCQCyeY,3804
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=AgVAdUbSkHXONVUjAyBQEXhIUUlinf9kNljcBpWnj3A,3276
-ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=nq94VpQ103eOimnmdyg7u3Xk1LH1IxGlmIbr2AttRIk,16224
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=TfbfsmuKoGsBENF9fYIAN_SMEQNhj-kjNdqQXFJGxpg,7784
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=L6hLaMQGb8-_BwSvTLIuDnZwfTqn0K4swBUjfPnYWZo,2341
 ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=FCbnwlkpYYb-tF7KscbSYjNEdg7XnuLju1cDuIRoQv8,8277
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=r9RqbyNvuvXOGu3ojtl7ZmbC7o4Pt8aUKAhN1yCdtEc,3397
@@ -56,22 +56,24 @@ ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py,sha256=5i
 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=q2gG5RRo7RgNzvHXYC0Juh6Tgt5d_RTMSWFaYvOKiZU,21065
-ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=anR99IrzR21x6yswFHYG5QQtPDZ7rVicf6STfMp54fU,8998
+ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=rRgwCEdVtzcJEaGbbBjw8HxCxrCX3pXA5nelawdYiME,9036
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=EV07_MEG3fv9g0ZGu9gbBd5BjjrGkxCT1pv7dvhz4TI,3791
+ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py,sha256=rzL5h7Z5DIEgfpc1pWgYHdKt2aR8ha_CUqTKQBSPBaU,5521
 ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=MUr6fSj2hBuYSlNbZtrBBpzqB_0WY-l_xYcd_TFFUjY,4831
 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=hVGpuI8gpj4Rn9k4otsRE22MSLFHBDlUOgioY6Ru6VI,5629
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=PxixRZb00v5BQkWbDwaJgke4Rd5LwzdWe0zH9SG4Tj0,9127
+ai_edge_torch/generative/layers/attention.py,sha256=zNIBXxCOA5Mz_F_dfBbKpIovhtcB6q5a-i8oAxls1d0,7071
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
-ai_edge_torch/generative/layers/builder.py,sha256=ZSBVLv5EOtCkSW_Z8C2Hd7jN52nIAA2as1-qpmHGbCg,3201
+ai_edge_torch/generative/layers/builder.py,sha256=WLTeDId9t3Xwt0h1zxzqoYyFvfrNzPKLskcl39q8Aqw,3403
 ai_edge_torch/generative/layers/feed_forward.py,sha256=4j2QaSCw59Jkk_ixKDpKEj7FLRauzuExTiSNRzAjAhE,2820
 ai_edge_torch/generative/layers/kv_cache.py,sha256=4uiZLO3om5G3--kT04Jt0esEYznbkJ7QLzSHfb8mjc4,3090
-ai_edge_torch/generative/layers/model_config.py,sha256=KpJRIHV5BJH8QOa7h6LXLZyC7UDWgbCEsw0CvArz49Q,4064
+ai_edge_torch/generative/layers/model_config.py,sha256=2zT9nyoyuuyk5ziiww0VSJ6_JO7pDf7uOYbO9O3OQc4,4249
 ai_edge_torch/generative/layers/normalization.py,sha256=M27eW3TcNK20oaXClXtfnu0lLWrAGrSKSsbegRWnj3c,1867
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=12SsCuoRuLNCwnFGe_pHDOZEBwBcqXs87Aj0PaWWw4E,1383
+ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=dYafGC205QE5CLIbBTCI-7eVvEGZEHzs1toPEhemeDs,3391
 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/quant_attrs.py,sha256=ffBALrrbrfiG_mrOr-f3B1Gc6PlAma9gtvVnfP7SDzI,1862
@@ -81,10 +83,10 @@ ai_edge_torch/generative/quantize/quant_recipes.py,sha256=CRA2ENevS-3usHqidWDe2w
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=OQ4ghQXknA1PPjuY-xBgAmOpaIBgYFM8F2YAIot06hE,1345
 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=1NfZxKo9Gx6CmVfd86K1FkmsNQnjzIV1ojBS85UGvT0,6500
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=i_SAW-hD8SaHuopMZI9IuXXDFn5uSTJa1nKZhaC3dAQ,6811
 ai_edge_torch/generative/test/test_quantize.py,sha256=f70sH1ZFzdCwYj0MG-eg54WOC4LasR0D8CTUYpjxZYM,3728
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/loader.py,sha256=c-ZOIDBVnat_5l2W5sWU7HQm7CL-wducS8poSu5PlUg,10107
+ai_edge_torch/generative/utilities/loader.py,sha256=r-_hSanSjLZ_YXFpZUb0Up94u5F8JHp70Vf2nlONPSg,11269
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=guDTv-12UUvJGl4eDvvZX3t4rRKewfXO8SpcYXM6gbc,16156
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
@@ -100,8 +102,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=ExThdTXqnWmGC3-F6sdXbXr8nYzkEe_qCz
 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.dev20240604.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/METADATA,sha256=0Hjm-wBxxZRbVdxcFhmOKVwsU52_l5F6IblrflOWMWk,1748
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/METADATA,sha256=2yFjQFvs93hoppwWWSJi-B9HC2n-h1s2rzjsyGXU1zI,1748
+ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/RECORD,,