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

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.int)
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
-  decode_token = torch.tensor([[0]], dtype=torch.int)
+  decode_token = torch.tensor([[0]], dtype=torch.long)
   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.int, device="cpu")
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, 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.int)
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
-  decode_token = torch.tensor([[0]], dtype=torch.int)
+  decode_token = torch.tensor([[0]], dtype=torch.long)
   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.int, device="cpu")
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, 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.IntTensor) -> torch.FloatTensor:
-    tokens = tokens.type(torch.int)
+  def forward(self, tokens: torch.LongTensor) -> torch.FloatTensor:
+    tokens = tokens.type(torch.long)
 
     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.int)
+  prompt_tokens = torch.full((1, n_tokens), 0, dtype=torch.long)
   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.int)
+  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
   prompt_e_token = [1, 2, 3, 4, 5, 6]
   prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(
-      prompt_e_token, dtype=torch.int
+      prompt_e_token, dtype=torch.long
   )
   prefill_e_input_pos = torch.arange(0, seq_len)
-  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
+  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
   prompt_d_token = [1, 2, 3, 4, 5, 6]
   prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(
-      prompt_d_token, dtype=torch.int
+      prompt_d_token, dtype=torch.long
   )
   prefill_d_input_pos = torch.arange(0, seq_len)
 
   # decoder
-  decode_token = torch.tensor([[1]], dtype=torch.int)
+  decode_token = torch.tensor([[1]], dtype=torch.long)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
-  decode_d_token = torch.tensor([[1]], dtype=torch.int)
+  decode_d_token = torch.tensor([[1]], dtype=torch.long)
   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.int)
+  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
   prompt_e_token = [1, 2, 3, 4, 5, 6]
   prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(
-      prompt_e_token, dtype=torch.int
+      prompt_e_token, dtype=torch.long
   )
   prefill_e_input_pos = torch.arange(0, seq_len)
-  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
+  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
   prompt_d_token = [1, 2, 3, 4, 5, 6]
   prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(
-      prompt_d_token, dtype=torch.int
+      prompt_d_token, dtype=torch.long
   )
   prefill_d_input_pos = torch.arange(0, seq_len)
 
   # decoder
-  decode_token = torch.tensor([[1]], dtype=torch.int)
+  decode_token = torch.tensor([[1]], dtype=torch.long)
   decode_input_pos = torch.tensor([0], dtype=torch.int64)
-  decode_d_token = torch.tensor([[1]], dtype=torch.int)
+  decode_d_token = torch.tensor([[1]], dtype=torch.long)
   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.int, device="cpu")
+  tokens = torch.full((1, 512), 0, dtype=torch.long, 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.int, device="cpu")
+  tokens = torch.full((1, 512), 0, dtype=torch.long, 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, dtype=torch.int), 0)
+  idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN), 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, dtype=torch.int), 0)
+  idx = torch.unsqueeze(torch.arange(0, 100), 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.int)
+  idx = torch.tensor([[1]], dtype=torch.long)
   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, dtype=torch.int), 0)
+  idx = torch.unsqueeze(torch.arange(0, 100), 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.int)
+  idx = torch.tensor([[1]], dtype=torch.long)
   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.int)
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
-  decode_token = torch.tensor([[0]], dtype=torch.int)
+  decode_token = torch.tensor([[0]], dtype=torch.long)
   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.int, device="cpu")
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, 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.int, device="cpu")
+  tokens = torch.full((1, 10), 0, dtype=torch.long, 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.int), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.long), 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.int), torch.tensor(
+    idx, input_pos = torch.tensor([[1], [2]], dtype=torch.long), 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.int), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.long), 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.int, device="cpu")
+    tokens = torch.full((1, 10), 0, dtype=torch.long, 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.int, device="cpu")
+    prefill_tokens = torch.full((1, seq_len), 0, dtype=torch.long, 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.int)
+    decode_token = torch.tensor([[1]], dtype=torch.long)
     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.int, device="cpu")
+    tokens = torch.full((1, 10), 0, dtype=torch.long, 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.int, device="cpu")
+    tokens = torch.full((1, 10), 0, dtype=torch.long, 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.int), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.long), 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.int), torch.tensor(
+    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
         [10], dtype=torch.int64
     )
 

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

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

@@ -36,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=yl36sMjVqDlf9I41DF9C5wx6ztMxYB5xukD1NltUS04,2536
-ai_edge_torch/generative/examples/gemma/gemma.py,sha256=caBQrJTK4tXFeGM-i2cNXd2Tb8GXi463MwFHc7N65WE,5934
+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/phi2/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-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/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/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=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/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/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
@@ -57,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=vqYip3JnjFMN5s0VnIzwhlBJl3up75WMb_VtRTfMOK0,4524
-ai_edge_torch/generative/examples/t5/t5.py,sha256=TyBDb50NbtKaHyGLGOMJ8dR2_GPF43oe8WhRfb1SMZ4,21085
+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/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=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/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/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-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/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/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
@@ -84,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=zgBgMyZ8RlSIjRhbaeodLyt6sj_dYuM7oI6Zyx7xQIg,1542
+ai_edge_torch/generative/quantize/example.py,sha256=t-YwyKSPAG-OZC1DfH-0vfie2RHHpTSQjxUY-tmhu5g,1543
 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
@@ -94,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=tXES8gePl8BptETyUgpIznbIK1SvGEd8mKz1bT_a8Mw,7581
-ai_edge_torch/generative/test/test_quantize.py,sha256=Emp-8oLHyGddELSCkncuSQt8ZJIhZ2-y0-ghR92s10g,5386
+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/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
@@ -114,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.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,,
+ai_edge_torch_nightly-0.2.0.dev20240702.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240702.dist-info/METADATA,sha256=yNbT2nyEGorcu8dwcB351bUhe0SLBaFWGhwOUBJE9zA,1748
+ai_edge_torch_nightly-0.2.0.dev20240702.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240702.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240702.dist-info/RECORD,,