ai-edge-torch-nightly 0.3.0.dev20240910__py3-none-any.whl → 0.3.0.dev20240911__py3-none-any.whl

ai_edge_torch/generative/layers/kv_cache.py

@@ -12,72 +12,181 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# `nn.Module` which implements a KV cache.
 
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+"""Utility functions for externalized KV Cache."""
+
+import dataclasses
+from typing import List, Tuple
+
+from ai_edge_torch import hlfb
+from ai_edge_torch.generative.layers import model_config
 import torch
-from torch import nn
+import torch.utils._pytree as pytree
 
 
-class KVCache(nn.Module):
+@dataclasses.dataclass
+class KVCacheEntry:
+  """A single cache entry that includes K and V caches.
 
-  def __init__(
-      self, batch_size, kv_cache_max, n_heads, head_dim, enable_hlfb=False
-  ):
-    """Initializes the KVCache layer.
+  The chaches are built based on the provided config with the shape of
+  (batch_size=1, kv_cache_max, num_query_groups, head_dim).
+  """
 
-    Args:
-      batch_size (int): batch size. Currently only batch size 1 is supported.
-      kv_cache_max (int): the max length of KV cache.
-      n_heads (int): number of kv heads.
-      head_dim (int): the head dimension size.
-      enable_hlfb (bool): whether hlfb is enabled or not.
-    """
-    super().__init__()
-    cache_shape = (batch_size, kv_cache_max, n_heads, head_dim)
-    self.register_buffer("k_cache", torch.zeros(cache_shape), persistent=False)
-    self.register_buffer("v_cache", torch.zeros(cache_shape), persistent=False)
-    self.enable_hlfb = enable_hlfb
-    self.kv_cache_max = kv_cache_max
+  k_cache: torch.Tensor
+  v_cache: torch.Tensor
 
-  def update_cache(self, input_pos, k_val, v_val):
-    """Update an entry in the KV cache.
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheEntry":
+    """Build an instance of the class based on model config."""
+    shape = (
+        1,  # Batch dimmension.
+        config.kv_cache_max,
+        config.attn_config.num_query_groups,
+        config.attn_config.head_dim,
+    )
+    k = torch.zeros(shape, dtype=dtype, device=device)
+    v = torch.zeros(shape, dtype=dtype, device=device)
+    obj = cls(k_cache=k, v_cache=v)
+    return obj
 
-    Args:
-      input_pos (torch.Tensor): the input position.
-      k_val (torch.Tensor): the new `key` value.
-      v_val (torch.Tensor): the new `value` value.
 
-    Returns:
-      The updated key and value tensor.
-    """
-    if self.enable_hlfb:
-      return self.update_cache_with_hlfb(input_pos, k_val, v_val)
+@dataclasses.dataclass
+class KVCache:
+  """A utility class for holding KV cache entries per layer."""
 
-    updated_k = self.k_cache.index_copy_(1, input_pos, k_val)
-    updated_v = self.v_cache.index_copy_(1, input_pos, v_val)
-    # Here we need a clone otherwise dynamo export will fail.
-    return torch.clone(updated_k), torch.clone(updated_v)
+  caches: Tuple[KVCacheEntry, ...]
 
-  def update_cache_with_hlfb(self, input_pos, k_val, v_val):
-    """Update an entry in the KV cache and enable high-level function boundary.
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCache":
+    """Build an instance of the class based on model config.
 
     Args:
-      input_pos (torch.Tensor): the input position.
-      k_val (torch.Tensor): the new `key` value.
-      v_val (torch.Tensor): the new `value` value.
+        config (ModelConfig): Model config used for building the cache.
+        dtype (torch.dtype, optional): The data type of the cache tensor.
+          Defaults to torch.float32.
+        device (torch.device, optional): The device placement of the cache
+          tensors. Defaults to None.
 
     Returns:
-      The updated key and value tensor.
+        KVCache: The created cache object.
     """
+    caches = [
+        KVCacheEntry.from_model_config(config, dtype, device)
+        for _ in range(config.num_layers)
+    ]
+    obj = cls(caches=tuple(caches))
+    return obj
 
-    builder = StableHLOCompositeBuilder(
-        name="odml.update_kv_cache", attr={"kv_cache_max": self.kv_cache_max}
-    )
-    k_cache, v_cache, input_pos, k_val, v_val = builder.mark_inputs(
-        self.k_cache, self.v_cache, input_pos, k_val, v_val
+  def flatten(self) -> List[torch.Tensor]:
+    """Flatten the cache entries into a list of tensors with order k_i, v_i."""
+    flattened, _ = _flatten_kvc(self)
+    return flattened
+
+
+def _flatten_kvc(kvc: KVCache) -> Tuple[List[str], List[str]]:
+  flattened = []
+  flat_names = []
+  none_names = []
+  for i, kv_entry in enumerate(kvc.caches):
+    flattened.append(kv_entry.k_cache)
+    flat_names.append(f"k_{i}")
+    flattened.append(kv_entry.v_cache)
+    flat_names.append(f"v_{i}")
+  return flattened, [flat_names, none_names]
+
+
+def _flatten_kvc_with_keys(kvc: KVCache) -> Tuple[List, List]:
+  flattened, (flat_names, none_names) = _flatten_kvc(kvc)
+  return [
+      (pytree.MappingKey(k), v) for k, v in zip(flat_names, flattened)
+  ], flat_names
+
+
+def _unflatten_kvc(
+    values: List[torch.Tensor], context: Tuple[List, List]
+) -> KVCache:
+  assert len(values) % 2 == 0, "Found odd number of K and V entries."
+  num_layers = len(values) // 2
+  flat_names = context[0]
+  kv_entries = []
+  for i in range(num_layers):
+    k_cache_idx = flat_names.index(f"k_{i}")
+    v_cache_idx = flat_names.index(f"v_{i}")
+    kv_entries.append(
+        KVCacheEntry(k_cache=values[k_cache_idx], v_cache=values[v_cache_idx])
     )
-    updated_k = k_cache.index_copy_(1, input_pos, k_val)
-    updated_v = v_cache.index_copy_(1, input_pos, v_val)
-    updated_k, updated_v = builder.mark_outputs(updated_k, updated_v)
-    return updated_k, updated_v
+  obj = KVCache(tuple(kv_entries))
+  return obj
+
+
+pytree.register_pytree_node(
+    KVCache,
+    _flatten_kvc,
+    _unflatten_kvc,
+    flatten_with_keys_fn=_flatten_kvc_with_keys,
+    serialized_type_name="",
+)
+
+
+def update(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+    enable_hlfb: bool = True,
+) -> KVCacheEntry:
+  """Out of place update of Cache buffer.
+
+  Args:
+      cache (KVCacheEntry): The original cache buffer.
+      input_pos (torch.Tensor): The update slice positions.
+      k_slice (torch.Tensor): The K slice to be updated in the new cache.
+      v_slice (torch.Tensor): The V slice to be updated in the new cache.
+      enable_hlfb (bool, optional): Whether the op is annotated for export with
+        High Level Function Boundary. Defaults to True.
+
+  Returns:
+      KVCacheEntry: The updated KVCache entry based on the passed inputs.
+  """
+  update_func = _update_kv_hlfb_impl if enable_hlfb else _update_kv_base_impl
+  return update_func(cache, input_pos, k_slice, v_slice)
+
+
+def _update_kv_base_impl(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntry:
+  """Update the cache buffer without High Level Function Boundary annotation."""
+  k = cache.k_cache.index_copy(1, input_pos, k_slice)
+  v = cache.v_cache.index_copy(1, input_pos, v_slice)
+  updated_cache = KVCacheEntry(k, v)
+  return updated_cache
+
+
+def _update_kv_hlfb_impl(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntry:
+  """Update the cache buffer with High Level Function Boundary annotation."""
+  builder = hlfb.StableHLOCompositeBuilder(name="odml.update_external_kv_cache")
+  k_cache, v_cache, input_pos, k_slice, v_slice = builder.mark_inputs(
+      cache.k_cache, cache.v_cache, input_pos, k_slice, v_slice
+  )
+  k = k_cache.index_copy(1, input_pos, k_slice)
+  v = v_cache.index_copy(1, input_pos, v_slice)
+  k, v = builder.mark_outputs(k, v)
+  return KVCacheEntry(k, v)

ai_edge_torch/generative/test/{test_experimental_ekv.py → test_kv_cache.py}

@@ -12,19 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# A suite of tests to validate experimental external KV Cache layers and models.
 
-from ai_edge_torch.generative.examples.experimental.gemma import gemma
-from ai_edge_torch.generative.examples.experimental.phi import phi2
-from ai_edge_torch.generative.examples.experimental.tiny_llama import tiny_llama  # NOQA
-from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
+"""A suite of tests to validate KV Cache layer."""
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import torch
 
 from absl.testing import absltest as googletest
 
 
-class TestExternalKVLayers(googletest.TestCase):
+class TestKVLayers(googletest.TestCase):
 
   def _get_test_config(
       self, num_layers, head_dim, num_query_groups, kv_cache_max_len
@@ -54,7 +52,7 @@ class TestExternalKVLayers(googletest.TestCase):
         num_query_groups=NUM_QG,
         kv_cache_max_len=KV_LEN,
     )
-    kv = kv_utils.EKVCache.from_model_config(config)
+    kv = kv_utils.KVCache.from_model_config(config)
     entry = kv.caches[0]
     # single-slice update
     input_pos = torch.tensor([1])
@@ -88,14 +86,14 @@ class TestExternalKVLayers(googletest.TestCase):
   def test_serialization(self):
     class TestModel(torch.nn.Module):
 
-      def forward(self, kv: kv_utils.EKVCache) -> kv_utils.EKVCache:
+      def forward(self, kv: kv_utils.KVCache) -> kv_utils.KVCache:
         updated_kv_entries = [
             kv_utils.KVCacheEntry(
                 torch.zeros_like(entry.k_cache), torch.zeros_like(entry.v_cache)
             )
             for entry in kv.caches
         ]
-        return kv_utils.EKVCache(updated_kv_entries)
+        return kv_utils.KVCache(updated_kv_entries)
 
     N = 1
     HEAD_DIM = 2
@@ -107,7 +105,7 @@ class TestExternalKVLayers(googletest.TestCase):
         num_query_groups=NUM_QG,
         kv_cache_max_len=KV_LEN,
     )
-    kv = kv_utils.EKVCache.from_model_config(config)
+    kv = kv_utils.KVCache.from_model_config(config)
     model = TestModel()
     exported_program = torch.export.export(model, (kv,))
     input_specs = exported_program.graph_signature.input_specs
@@ -116,17 +114,5 @@ class TestExternalKVLayers(googletest.TestCase):
     self.assertEqual(input_specs[1].arg.name, "kv_v_0")
 
 
-class TestExternalKVModels(googletest.TestCase):
-
-  def test_can_build_gemma(self):
-    gemma.define_and_run_2b(checkpoint_path=None, test_model=True)
-
-  def test_can_build_phi2(self):
-    phi2.define_and_run(checkpoint_path=None, test_model=True)
-
-  def test_can_build_tinyllama(self):
-    tiny_llama.define_and_run(checkpoint_path=None, test_model=True)
-
-
 if __name__ == "__main__":
   googletest.main()

ai_edge_torch/generative/test/test_model_conversion.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Testing model conversion for a few gen-ai models.
-import copy
+
+"""Testing model conversion for a few gen-ai models."""
 
 import ai_edge_torch
 from ai_edge_torch import config as ai_edge_config
-from ai_edge_torch.generative.examples.gemma import gemma, gemma2
-from ai_edge_torch.generative.examples.phi2 import phi2
-from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache  # NOQA
+from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
-from ai_edge_torch.testing import model_coverage
+from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.test import utils as test_utils
 import numpy as np
 import torch
 
@@ -49,22 +48,32 @@ class TestModelConversion(googletest.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).eval()
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKVCache(config).eval()
+    tokens, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
         [10], dtype=torch.int64
     )
-
-    edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (idx, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-5,
             rtol=1e-5,
         )
@@ -77,22 +86,32 @@ class TestModelConversion(googletest.TestCase):
   def test_toy_model_with_kv_cache_with_hlfb(self):
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True
-    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config).eval()
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKVCache(config).eval()
+    tokens, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
         [10], dtype=torch.int64
     )
-
-    edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (idx, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-5,
             rtol=1e-5,
         )
@@ -117,37 +136,56 @@ class TestModelConversion(googletest.TestCase):
     decode_token = torch.tensor([[1]], dtype=torch.long)
     decode_input_pos = torch.tensor([5], dtype=torch.int64)
 
+    kv = kv_cache.KVCache.from_model_config(config)
+
     edge_model = (
         ai_edge_torch.signature(
-            "prefill", pytorch_model, (prefill_tokens, prefill_input_pos)
+            "prefill",
+            pytorch_model,
+            sample_kwargs={
+                "tokens": prefill_tokens,
+                "input_pos": prefill_input_pos,
+                "kv_cache": kv,
+            },
+        )
+        .signature(
+            "decode",
+            pytorch_model,
+            sample_kwargs={
+                "tokens": decode_token,
+                "input_pos": decode_input_pos,
+                "kv_cache": kv,
+            },
         )
-        .signature("decode", pytorch_model, (decode_token, decode_input_pos))
         .convert()
     )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
-    copied_model = copy.deepcopy(pytorch_model)
-    copied_edge = copy.deepcopy(edge_model)
-
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (prefill_tokens, prefill_input_pos),
+            prefill_tokens,
+            prefill_input_pos,
+            kv,
             signature_name="prefill",
-            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
         )
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            copied_edge,
-            copied_model,
-            (decode_token, decode_input_pos),
+        test_utils.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            decode_token,
+            decode_input_pos,
+            kv,
             signature_name="decode",
-            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
         )
     )
 

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Testing model conversion for a few gen-ai models.
-import copy
+
+"""Testing model conversion for a few gen-ai models."""
 
 import ai_edge_torch
 from ai_edge_torch import config as ai_edge_config
-from ai_edge_torch.generative.examples.gemma import gemma, gemma2
-from ai_edge_torch.generative.examples.phi2 import phi2
-from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache  # NOQA
-from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
-from ai_edge_torch.testing import model_coverage
+from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.examples.phi import phi2
+from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.test import utils as test_utils
 import numpy as np
 import torch
 
@@ -55,18 +55,28 @@ class TestModelConversion(googletest.TestCase):
     tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
     tokens[0, :4] = idx
     input_pos = torch.arange(0, 10)
-
-    edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-2,
             rtol=1e-5,
         )
@@ -85,23 +95,31 @@ class TestModelConversion(googletest.TestCase):
     prefill_tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
     prefill_tokens[0, :4] = idx
     prefill_input_pos = torch.arange(0, 10)
+    kv = kv_cache.KVCache.from_model_config(config)
 
     edge_model = ai_edge_torch.signature(
-        "prefill", model, (prefill_tokens, prefill_input_pos)
+        "prefill",
+        model,
+        sample_kwargs={
+            "tokens": prefill_tokens,
+            "input_pos": prefill_input_pos,
+            "kv_cache": kv,
+        },
     ).convert()
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             model,
-            (prefill_tokens, prefill_input_pos),
+            prefill_tokens,
+            prefill_input_pos,
+            kv,
             signature_name="prefill",
-            num_valid_inputs=1,
-            atol=1e-2,
-            rtol=1e-5,
+            atol=1e-1,
+            rtol=1e-3,
         )
     )
 
@@ -117,18 +135,28 @@ class TestModelConversion(googletest.TestCase):
     tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
     tokens[0, :4] = idx
     input_pos = torch.arange(0, 10)
-
-    edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-3,
             rtol=1e-3,
         )

ai_edge_torch/generative/test/utils.py

@@ -0,0 +1,54 @@
+# 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.
+# ==============================================================================
+
+"""Common utils for testing."""
+
+from ai_edge_torch import model
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.lowertools import common_utils
+import numpy as np
+import torch
+from torch.utils import _pytree as pytree
+
+
+def compare_tflite_torch(
+    edge_model: model.Model,
+    torch_model: torch.nn.Module,
+    tokens: torch.Tensor,
+    input_pos: torch.Tensor,
+    kv_cache: kv_utils.KVCache,
+    signature_name: str,
+    atol: float = 1e-5,
+    rtol: float = 1e-5,
+):
+  """Compares torch models and TFLite models."""
+  values, spec = pytree.tree_flatten({"kv_cache": kv_cache})
+  flat_names = common_utils.flat_dict_names(spec.children_specs, spec.context)
+  torch_output = torch_model(tokens, input_pos, kv_cache)
+
+  input_kv_flatten = {k: v.numpy() for k, v in zip(flat_names, values)}
+  edge_output = edge_model(
+      signature_name=signature_name,
+      tokens=tokens.numpy(),
+      input_pos=input_pos.numpy(),
+      **input_kv_flatten,
+  )
+
+  return np.allclose(
+      edge_output["logits"],
+      torch_output["logits"].detach().numpy(),
+      atol=atol,
+      rtol=rtol,
+  )

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20240910"
+__version__ = "0.3.0.dev20240911"

{ai_edge_torch_nightly-0.3.0.dev20240910.dist-info → ai_edge_torch_nightly-0.3.0.dev20240911.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240910
+Version: 0.3.0.dev20240911
 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