ai-edge-torch-nightly 0.2.0.dev20240605__py3-none-any.whl → 0.2.0.dev20240607__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_nightly-0.2.0.dev20240605.dist-info → ai_edge_torch_nightly-0.2.0.dev20240607.dist-info}/METADATA

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

@@ -83,7 +83,7 @@ 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=r-_hSanSjLZ_YXFpZUb0Up94u5F8JHp70Vf2nlONPSg,11269
@@ -102,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.dev20240605.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/METADATA,sha256=GJzwmKkM4T0H-vTvMyoxiD80WfppEpE_sd2Ip4aSbgM,1748
-ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240607.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240607.dist-info/METADATA,sha256=wXfrp5POJLq610NgMNlMi-WaC2AHCc0u9d4z2t6KiMg,1748
+ai_edge_torch_nightly-0.2.0.dev20240607.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240607.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240607.dist-info/RECORD,,