ai-edge-torch-nightly 0.7.0.dev20251026__py3-none-any.whl → 0.7.0.dev20251028__py3-none-any.whl

ai_edge_torch/generative/test/test_kv_cache.py

@@ -41,6 +41,20 @@ class TestKVLayers(googletest.TestCase):
     )
     return config
 
+  def _assert_kv_cache_entry_equal(self, kv1, kv2):
+    self.assertIsInstance(kv1, kv_utils.KVCacheEntry)
+    self.assertIsInstance(kv2, kv_utils.KVCacheEntry)
+    self.assertEqual(kv1.kv_layout, kv2.kv_layout)
+    self.assertTrue(torch.equal(kv1.k_cache, kv2.k_cache))
+    self.assertTrue(torch.equal(kv1.v_cache, kv2.v_cache))
+
+  def _assert_kv_cache_equal(self, kv1, kv2):
+    self.assertIsInstance(kv1, kv_utils.KVCache)
+    self.assertIsInstance(kv2, kv_utils.KVCache)
+    self.assertEqual(len(kv1.caches), len(kv2.caches))
+    for kv1_entry, kv2_entry in zip(kv1.caches, kv2.caches):
+      self._assert_kv_cache_entry_equal(kv1_entry, kv2_entry)
+
   def test_cache_udpate(self):
     N = 1
     HEAD_DIM = 2
@@ -118,7 +132,7 @@ class TestKVLayers(googletest.TestCase):
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, NUM_LAYERS * 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
-    self.assertEqual(kv, kv_unflat)
+    self._assert_kv_cache_equal(kv, kv_unflat)
 
   def test_pytree_roundtrip_kv_cache_derived(self):
     NUM_LAYERS = 4
@@ -134,7 +148,7 @@ class TestKVLayers(googletest.TestCase):
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, NUM_LAYERS * 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
-    self.assertEqual(kv, kv_unflat)
+    self._assert_kv_cache_equal(kv, kv_unflat)
 
   def test_pytree_roundtrip_kv_entry(self):
     attn_config = cfg.AttentionConfig(
@@ -144,8 +158,7 @@ class TestKVLayers(googletest.TestCase):
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
-    self.assertEqual(kv, kv_unflat)
-    self.assertIsInstance(kv_unflat, kv_utils.KVCacheEntry)
+    self._assert_kv_cache_entry_equal(kv, kv_unflat)
 
   def test_pytree_roundtrip_kv_entry_derived(self):
     attn_config = cfg.AttentionConfig(
@@ -157,8 +170,7 @@ class TestKVLayers(googletest.TestCase):
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
-    self.assertEqual(kv, kv_unflat)
-    self.assertIsInstance(kv_unflat, kv_utils.KVCacheEntry)
+    self._assert_kv_cache_entry_equal(kv, kv_unflat)
 
 
 if __name__ == "__main__":

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.7.0.dev20251026"
+__version__ = "0.7.0.dev20251028"

{ai_edge_torch_nightly-0.7.0.dev20251026.dist-info → ai_edge_torch_nightly-0.7.0.dev20251028.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=J23twq79-lJ3kkP65CX-o5VDaUH2L3KosiRYZcyBkpo,806
+ai_edge_torch/version.py,sha256=nyYzWr9H5qeB7EwfB29uEskmVnED_E-YkjvbTfteADc,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -200,7 +200,7 @@ ai_edge_torch/generative/quantize/quant_recipes.py,sha256=mmPsgQ2vloMWflqJ6ALmD1
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=TwR2FpQuBEORy6FshEyHNBMKARWlA2MVtTfX9tXV5aE,1488
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
-ai_edge_torch/generative/test/test_kv_cache.py,sha256=zVocQZTORIROwSYVMmy5KGPDl6kTuEGko7L1PTkua8g,5568
+ai_edge_torch/generative/test/test_kv_cache.py,sha256=kdoBpyzVOjQsZkEPnbDOh33gRE0RR8tq1XCgJkSYwU0,6157
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=sKnBixmGIXHUtOwh3SA4MIFenPbjK2n-Xknwic_KMDQ,5046
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=T35zdzag2-nmy4qc6AifAjbDXAHU2vyLTE1QCabYBzk,6298
@@ -270,8 +270,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.7.0.dev20251026.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20251026.dist-info/METADATA,sha256=APrzqTt83_zoXvSYgUBafxFb2clguJzGpfSd3Hf0ajg,2074
-ai_edge_torch_nightly-0.7.0.dev20251026.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20251026.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20251026.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20251028.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20251028.dist-info/METADATA,sha256=l8GlfrlrRJuJIeQMp1rCHPjaOBXjCND_vZbgAg37iVg,2074
+ai_edge_torch_nightly-0.7.0.dev20251028.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.7.0.dev20251028.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20251028.dist-info/RECORD,,