ai-edge-torch-nightly 0.2.0.dev20240711__py3-none-any.whl → 0.2.0.dev20240712__py3-none-any.whl

ai_edge_torch/convert/conversion.py

@@ -88,16 +88,14 @@ def convert_signatures(
   _warn_training_modules(signatures)
 
   exported_programs: torch.export.ExportedProgram = [
-      torch.export.export(
-          sig.module, sig.sample_args, dynamic_shapes=sig.dynamic_shapes
-      )
+      torch.export.export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)
       for sig in signatures
   ]
 
   # Apply default fx passes
   exported_programs = list(map(_run_convert_passes, exported_programs))
   shlo_bundles: list[stablehlo.StableHLOModelBundle] = [
-      cutils.exported_program_to_stablehlo_bundle(exported, sig.sample_args)
+      cutils.exported_program_to_stablehlo_bundle(exported, sig.flat_args)
       for exported, sig in zip(exported_programs, signatures)
   ]
 

ai_edge_torch/convert/conversion_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import collections
 import copy
 from dataclasses import dataclass
 import gc
@@ -22,6 +23,7 @@ import tempfile
 from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
+import torch.utils._pytree as pytree
 from torch_xla import stablehlo
 
 from ai_edge_torch.generative.quantize.ai_edge_quantizer_glue import translate_recipe  # NOQA
@@ -47,8 +49,59 @@ class Signature:
   name: str
   module: torch.nn.Module
   sample_args: tuple[torch.Tensor]
+  sample_kwargs: dict[str, torch.Tensor]
   dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None
 
+  @property
+  def _normalized_sample_args_kwargs(self):
+    args, kwargs = self.sample_args, self.sample_kwargs
+    if args is not None:
+      if not isinstance(args, tuple):
+        # TODO(b/352584188): Check value types
+        raise ValueError("sample_args must be a tuple of torch tensors.")
+    if kwargs is not None:
+      if not isinstance(kwargs, dict) or not all(
+          isinstance(key, str) for key in kwargs.keys()
+      ):
+        # TODO(b/352584188): Check value types
+        raise ValueError("sample_kwargs must be a dict of string to tensor.")
+
+    args = args if args is not None else tuple()
+    kwargs = kwargs if kwargs is not None else {}
+    return args, kwargs
+
+  @property
+  def flat_arg_names(self) -> list[str]:
+    spec = pytree.tree_flatten(self._normalized_sample_args_kwargs)[1]
+    args_spec, kwargs_spec = spec.children_specs
+
+    names = []
+    for i in range(args_spec.num_leaves):
+      names.append(f"args_{i}")
+
+    dict_context = (
+        kwargs_spec.context
+        if kwargs_spec.type is not collections.defaultdict
+        # ignore mismatch of `default_factory` for defaultdict
+        else kwargs_spec.context[1]
+    )
+
+    for name, value_spec in zip(dict_context, kwargs_spec.children_specs):
+      if value_spec.num_leaves == 1:
+        names.append(name)
+      else:
+        # value_spec.num_leaves may be greater than 1 when the value is a (nested)
+        # tuple of tensors. We haven't decided how we should support flattenable
+        # tensor containers as inputs.
+        # TODO(b/352584188): Decide the behavior of tensor container as input (flatten or reject)
+        for i in range(value_spec.num_leaves):
+          names.append(f"{name}_{i}")
+    return names
+
+  @property
+  def flat_args(self) -> tuple[torch.Tensor]:
+    return tuple(pytree.tree_flatten(self._normalized_sample_args_kwargs)[0])
+
 
 def exported_program_to_stablehlo_bundle(
     exported_program: torch.export.ExportedProgram, sample_args: tuple[torch.Tensor]
@@ -189,7 +242,9 @@ def _make_tf_function(
 
 def _make_tf_signature(
     meta: stablehlo.StableHLOFunctionMeta,
+    signature: Signature,
 ) -> list[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
   input_pos_to_spec = {
       loc.position: spec
       for loc, spec in itertools.chain(
@@ -197,9 +252,11 @@ def _make_tf_signature(
       )
       if loc.type_ == stablehlo.VariableType.INPUT_ARG
   }
+  assert len(input_pos_to_spec) == len(input_names)
+
   primitive_type_to_tf_type = {"int": "int32", "float": "float32"}
   ret: list[tf.TensorSpec] = []
-  for i in range(len(input_pos_to_spec)):
+  for i, name in enumerate(input_names):
     spec = input_pos_to_spec[i]
     shape = _get_shape_with_dynamic(spec)
     ret.append(
@@ -208,7 +265,7 @@ def _make_tf_signature(
             dtype=primitive_type_to_tf_type[spec.dtype]
             if spec.dtype in primitive_type_to_tf_type
             else spec.dtype,
-            name=f"args_{i}",
+            name=name,
         )
     )
   return ret
@@ -276,7 +333,8 @@ def convert_stablehlo_to_tflite(
       tf.Variable(v, trainable=False) for v in bundle.additional_constants
   ]
   tf_signatures: list[list[tf.TensorSpec]] = list(
-      _make_tf_signature(func.meta) for func in bundle.stablehlo_funcs
+      _make_tf_signature(func.meta, sig)
+      for func, sig in zip(bundle.stablehlo_funcs, signatures)
   )
 
   tf_functions = _make_tf_function(shlo_graph_module, bundle)

ai_edge_torch/convert/converter.py

@@ -34,17 +34,23 @@ class Converter:
       self,
       name: str,
       module: torch.nn.Module,
-      sample_args: tuple[cutils.TracingArg],
+      sample_args=None,
+      sample_kwargs=None,
+      *,
       dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
   ) -> Converter:
     """Alias to `add_signature`"""
-    return self.add_signature(name, module, sample_args, dynamic_shapes)
+    return self.add_signature(
+        name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+    )
 
   def add_signature(
       self,
       name: str,
       module: torch.nn.Module,
-      sample_args: tuple[cutils.TracingArg],
+      sample_args=None,
+      sample_kwargs=None,
+      *,
       dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
   ) -> Converter:
     """Allows adding a new named torch model along with sample args to the conversion.
@@ -52,7 +58,8 @@ class Converter:
     Args:
       name: The name of the signature included in the converted edge model.
       module: The torch module to be converted.
-      sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
       dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
         See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
 
@@ -63,13 +70,21 @@ class Converter:
     if name in [sig.name for sig in self._signatures]:
       raise ValueError(f"A signature with the provided name ({name}) is already added.")
 
-    self._signatures.append(cutils.Signature(name, module, sample_args, dynamic_shapes))
+    if sample_args is None and sample_kwargs is None:
+      raise ValueError("sample_args or sample_kwargs must be provided.")
+
+    self._signatures.append(
+        cutils.Signature(
+            name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+        )
+    )
     return self
 
   def convert(
       self,
       module: torch.nn.Module = None,
-      sample_args: tuple[cutils.TracingArg] = None,
+      sample_args=None,
+      sample_kwargs=None,
       *,
       quant_config: Optional[qcfg.QuantConfig] = None,
       dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
@@ -88,7 +103,8 @@ class Converter:
     Args:
       name: The name of the signature included in the converted edge model.
       module: The torch module to be converted.
-      sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
       quant_config: User-defined quantization method and scheme of the model.
       dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
         See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
@@ -100,12 +116,20 @@ class Converter:
       ValueError: If the arguments are not provided as expected. See the example in this functions's comment.
     """
     if module is not None:
-      if sample_args is not None:  # both module and args provided
+      if (
+          sample_args is not None or sample_kwargs is not None
+      ):  # both module and args provided
         self.add_signature(
-            cutils.DEFAULT_SIGNATURE_NAME, module, sample_args, dynamic_shapes
+            cutils.DEFAULT_SIGNATURE_NAME,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+      else:  # module is provided but not args
+        raise ValueError(
+            "sample_args or sample_kwargs must be provided if a module is specified."
         )
-      else:  # module is provided but not sample_args
-        raise ValueError("sample_args needs to be provided if a module is specified.")
 
     return conversion.convert_signatures(
         self._signatures,
@@ -117,7 +141,8 @@ class Converter:
 def signature(
     name: str,
     module: torch.nn.Module,
-    sample_args: tuple[cutils.TracingArg],
+    sample_args=None,
+    sample_kwargs=None,
     dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
 ) -> Converter:
   """Initiates a Converter object with the provided signature.
@@ -125,7 +150,8 @@ def signature(
   Args:
     name: The name of the signature included in the converted edge model.
     module: The torch module to be converted.
-    sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
     dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
       See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
 
@@ -134,12 +160,15 @@ def signature(
     edge_model = converter.convert()
 
   """
-  return Converter().signature(name, module, sample_args, dynamic_shapes)
+  return Converter().signature(
+      name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+  )
 
 
 def convert(
     module: torch.nn.Module = None,
-    sample_args: tuple[cutils.TracingArg] = None,
+    sample_args=None,
+    sample_kwargs=None,
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
@@ -149,7 +178,8 @@ def convert(
 
   Args:
     module: The torch module to be converted.
-    sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
     quant_config: User-defined quantization method and scheme of the model.
     dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
       See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
@@ -165,6 +195,7 @@ def convert(
   return Converter().convert(
       module,
       sample_args,
+      sample_kwargs,
       quant_config=quant_config,
       dynamic_shapes=dynamic_shapes,
       _ai_edge_converter_flags=_ai_edge_converter_flags,

ai_edge_torch/convert/test/test_convert.py

@@ -267,6 +267,45 @@ class TestConvert(unittest.TestCase):
           model_coverage.compare_tflite_torch(edge_model, model, validate_input)
       )
 
+  def test_convert_model_with_kwargs(self):
+    """
+    Test converting a simple model with sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    kwargs_gen = lambda: dict(x=torch.randn(10, 10), y=torch.randn(10, 10))
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(model, sample_kwargs=kwargs_gen())
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, model, kwargs=kwargs_gen)
+    )
+
+  def test_convert_model_with_args_kwargs(self):
+    """
+    Test converting a simple model with both sample_args and sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    args_gen = lambda: (torch.randn(10, 10),)
+    kwargs_gen = lambda: dict(y=torch.randn(10, 10))
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(model, args_gen(), kwargs_gen())
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, model, args_gen, kwargs_gen)
+    )
+
 
 if __name__ == "__main__":
   unittest.main()

ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py

@@ -40,7 +40,7 @@ class RemoveSDPACompositeZeroMaskPass(ExportedProgramPassBase):
         if self.is_zero_tensor_node(source):
           # Remove the mark_tensor call on the mask input by
           # replacing the target with an identity function.
-          node.target = lambda *args, **kwargs: args[0]
+          node.target = lambda *args, **kwargs: torch.zeros_like(args[0])
 
     exported_program.graph_module.graph.lint()
     exported_program.graph_module.recompile()

ai_edge_torch/model.py

@@ -33,7 +33,10 @@ class Model(abc.ABC):
 
   @abc.abstractmethod
   def __call__(
-      self, *args: npt.ArrayLike, signature_name: str = cutils.DEFAULT_SIGNATURE_NAME
+      self,
+      *args: npt.ArrayLike,
+      signature_name: str = cutils.DEFAULT_SIGNATURE_NAME,
+      **kwargs,
   ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
     raise NotImplementedError()
 
@@ -62,12 +65,16 @@ class TfLiteModel(Model):
     self._tflite_model = tflite_model
 
   def __call__(
-      self, *args: npt.ArrayLike, signature_name: str = cutils.DEFAULT_SIGNATURE_NAME
+      self,
+      *args: npt.ArrayLike,
+      signature_name: str = cutils.DEFAULT_SIGNATURE_NAME,
+      **kwargs,
   ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
     """Runs inference on the edge model using the provided arguments.
 
     Args:
       *args: The arguments to be passed to the model for inference.
+      **kwargs: The arguments with specific names to be passed to the model for inference.
       signature_name: The name of the signature to be used for inference.
         The default signature is used if not provided.
     """
@@ -90,13 +97,14 @@ class TfLiteModel(Model):
       else:
         raise exception
 
-    if len(signature_list[signature_name]['inputs']) != len(args):
+    if len(signature_list[signature_name]['inputs']) != len(args) + len(kwargs):
       raise ValueError(
           f"The model requires {len(signature_list[signature_name]['inputs'])} arguments but {len(args)} was provided."
       )
 
     # Gather the input dictionary based on the signature.
     inputs = {f'args_{idx}': args[idx] for idx in range(len(args))}
+    inputs = {**inputs, **kwargs}
     outputs = runner(**inputs)
 
     return (

ai_edge_torch/testing/model_coverage/model_coverage.py

@@ -60,7 +60,8 @@ def _torch_tensors_to_np(*argv):
 def compare_tflite_torch(
     edge_model: Model,
     torch_eval_func: Callable,
-    input_data=None,
+    args=None,
+    kwargs=None,
     *,
     num_valid_inputs: int = 1,
     signature_name: str = None,
@@ -71,8 +72,9 @@ def compare_tflite_torch(
   Args:
     edge_model: Serialized ai_edge_torch.model.Model object.
     torch_eval_func: Callable function to evaluate torch model.
-    input_data: torch.tensor array  or a callable to generate a torch.tensor array
+    args: torch.tensor array or a callable to generate a torch.tensor array
       with random data, to pass into models during inference. (default None).
+    kwargs: dict of str to torch.tensor, or a callable to generate such.
     num_valid_inputs: Defines the number of times the random inputs will be generated (if a callable is provided for input_data).
     signature_name: If provided, specifies the name for the signature of the edge_model to run.
       Calls the default signature if not provided.
@@ -86,29 +88,33 @@ def compare_tflite_torch(
   # The supplied model_def.forward_args() will be executed num_valid_inputs
   # times to generate num_valid_inputs random inputs.
   torch_inputs = [
-      input_data() if callable(input_data) else input_data
+      (
+          (args() if callable(args) else args) or tuple(),
+          (kwargs() if callable(kwargs) else kwargs) or {},
+      )
       for _ in range(num_valid_inputs)
   ]
-  torch_outputs = [torch_eval_func(*xs) for xs in torch_inputs]
-  np_inputs = [_torch_tensors_to_np(xs) for xs in torch_inputs]
+  torch_outputs = [torch_eval_func(*args, **kwargs) for args, kwargs in torch_inputs]
+  np_inputs = [
+      (_torch_tensors_to_np(args), _torch_tensors_to_np(kwargs))
+      for args, kwargs in torch_inputs
+  ]
   np_outputs = [_torch_tensors_to_np(_flatten(ys)) for ys in torch_outputs]
 
   # Define inline utility function used throughout the function.
   def equal_fn(actual, expected):
     return np.allclose(actual, expected, atol=atol, rtol=rtol)
 
-  def get_actual_fn(input):
+  def get_edge_output(inputs):
+    args, kwargs = inputs
     if signature_name is None:
-      return _flatten(edge_model(*input))
+      return _flatten(edge_model(*args, **kwargs))
     else:
-      return _flatten(edge_model(*input, signature_name=signature_name))
-
-  def get_expected_fn(input=None, idx=0):
-    return np_outputs[idx]
+      return _flatten(edge_model(*args, **kwargs, signature_name=signature_name))
 
   for idx, np_input in enumerate(np_inputs):
-    output = get_actual_fn(np_input)
-    golden_output = get_expected_fn(np_input, idx)
+    output = get_edge_output(np_input)
+    golden_output = np_outputs[idx]
 
     is_output_len_eq = len(golden_output) == len(output)
 

{ai_edge_torch_nightly-0.2.0.dev20240711.dist-info → ai_edge_torch_nightly-0.2.0.dev20240712.dist-info}/METADATA

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

@@ -1,9 +1,9 @@
 ai_edge_torch/__init__.py,sha256=CNDboRP4zQBpz2hznNCQWcQCARvNXUm3DMa1Dw_XXFg,1067
-ai_edge_torch/model.py,sha256=kmcgELjsYl8YzF8nUF6P7q4i8MWS-pLGpfsy-yTUXmE,4243
+ai_edge_torch/model.py,sha256=8Ba9ia7TCM_fciulw6qObmzdcxL3IaLQKDqpR7Lxp-Q,4440
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/convert/conversion.py,sha256=8K8jQuaCjlUWoj7jiimxp_zpN6mYThLOcQ858UDcYnE,4159
-ai_edge_torch/convert/conversion_utils.py,sha256=9BqCL38DErv1vEVGtT3BIJVhdwZjw2EQ-_m5UpvVVYE,11294
-ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
+ai_edge_torch/convert/conversion.py,sha256=StJHglvx6cii36oi8sj-tZda009e9UqR6ufZOZkP1SY,4137
+ai_edge_torch/convert/conversion_utils.py,sha256=PKXIlSCU-8DhppNBh9ICDNUlEOpV0HgCbt85jDVe3rA,13394
+ai_edge_torch/convert/converter.py,sha256=hSrW6A-kix9cjdD6CuLL7rseWrLKoV6GRy-iUSW_nZc,7875
 ai_edge_torch/convert/to_channel_last_io.py,sha256=zo5tY3yDhY_EPCkrL1XSXs2uRFS8B4_qu08dSjNsUGk,2778
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=ijVyDclPnd6a0DWWUJkwR4igj6f82S-cE1-83QGPvgw,1652
@@ -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/greedy.py,sha256=8uHJbIwPMTgeSfYVba163pkXSQkHLxFwar_8A1AhgAM,2279
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=lklGxE1R32vsjFbhLLBDEFL4pfLi_iTgI9Ftb6Grezk,7156
 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.py,sha256=h0vOffr8saDQRkiXljNWDZ17EBjnS4xAtxd8DxETleY,9081
 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/convert/test/test_to_channel_last_io.py,sha256=I8c4ZG3v1vo0yxQYzLK_BTId4AOL9vadHGDtfCUZ4UI,2930
@@ -68,7 +68,7 @@ ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TI
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=nT7Fh-f5ZdwaK3dPoCvZflpJ4fRHjLdFMjk1_uw3-b8,2559
 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/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=BCAcc_OcEjvbaXQSbc8vlKeMad7E3gCA4BNsUdWRwBI,1966
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=AW0Qo3uOIe6p1rJNJ6zR_r4fqL2y-6QJHh0yUd-5Yb0,11966
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
@@ -113,9 +113,9 @@ ai_edge_torch/quantize/pt2e_quantizer_utils.py,sha256=yjzKoptnfEeW_sN7sODUfj3nCt
 ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDdN5XtvHwjc,3148
 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.dev20240711.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240711.dist-info/METADATA,sha256=GftPz7zSGYCaTvO4gntWftMbj0NCSh4OXJEe1epdBCU,1745
-ai_edge_torch_nightly-0.2.0.dev20240711.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240711.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240711.dist-info/RECORD,,
+ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=kzIulTldq8R9E-lAZsvfSTvLu3FYEX7b9DyYM3qisXM,4485
+ai_edge_torch_nightly-0.2.0.dev20240712.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240712.dist-info/METADATA,sha256=BGHmRLYo3ko7KRysDP59YexTpPn45jrtpRwqQkPAM5s,1745
+ai_edge_torch_nightly-0.2.0.dev20240712.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240712.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240712.dist-info/RECORD,,