ai-edge-torch-nightly 0.2.0.dev20240718__py3-none-any.whl → 0.2.0.dev20240719__py3-none-any.whl

ai_edge_torch/convert/conversion_utils.py

@@ -20,7 +20,7 @@ import gc
 import itertools
 import logging
 import tempfile
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.utils._pytree as pytree
@@ -79,28 +79,49 @@ class Signature:
     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]
+    kwargs_names = self._flat_kwarg_names(
+        kwargs_spec.children_specs, kwargs_spec.context
     )
+    names.extend(kwargs_names)
+    return names
 
-    for name, value_spec in zip(dict_context, kwargs_spec.children_specs):
-      if value_spec.num_leaves == 1:
-        names.append(name)
+  def _flat_kwarg_names(self, specs, context) -> List[str]:
+    flat_names = []
+    if context is None:
+      for i, spec in enumerate(specs):
+        if spec.children_specs:
+          flat_names.extend(
+              [
+                  f"{i}_{name}"
+                  for name in self._flat_kwarg_names(spec.children_specs, spec.context)
+              ]
+          )
+        else:
+          flat_names.append(f"{i}")
+    else:
+      flat_ctx = self._flatten_list(context)
+      for prefix, spec in zip(flat_ctx, specs):
+        leaf_flat_names = self._flat_kwarg_names(spec.children_specs, spec.context)
+        if leaf_flat_names:
+          flat_names.extend([f"{prefix}_{name}" for name in leaf_flat_names])
+        else:
+          flat_names.append(prefix)
+
+    return flat_names
+
+  def _flatten_list(self, l: List) -> List:
+    flattened = []
+    for item in l:
+      if isinstance(item, list):
+        flattened.extend(self._flatten_list(item))
       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
+        flattened.append(item)
+    return flattened
 
   @property
-  def flat_args(self) -> tuple[torch.Tensor]:
-    return tuple(pytree.tree_flatten(self._normalized_sample_args_kwargs)[0])
+  def flat_args(self) -> tuple[Any]:
+    args, kwargs = self._normalized_sample_args_kwargs
+    return tuple([*args, *kwargs.values()])
 
 
 def exported_program_to_stablehlo_bundle(

ai_edge_torch/convert/test/test_convert.py

@@ -14,10 +14,14 @@
 # ==============================================================================
 
 
+from dataclasses import dataclass
 import os
 import tempfile
+from typing import Tuple
 import unittest
 
+import numpy as np
+import tensorflow as tf
 import torch
 import torchvision
 
@@ -26,6 +30,15 @@ from ai_edge_torch.convert import conversion_utils as cutils
 from ai_edge_torch.testing import model_coverage
 
 
+@dataclass
+class TestContainer1:
+  data_1: torch.Tensor
+  data_2: Tuple[torch.Tensor, torch.Tensor]
+
+
+torch.export.register_dataclass(TestContainer1, serialized_type_name="TestContainer1")
+
+
 class TestConvert(unittest.TestCase):
   """Tests conversion of various modules."""
 
@@ -306,6 +319,99 @@ class TestConvert(unittest.TestCase):
         model_coverage.compare_tflite_torch(edge_model, model, args_gen, kwargs_gen)
     )
 
+  def test_convert_model_with_args_nested_kwargs_1(self):
+    """
+    Test converting a simple model with both sample_args and nested sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x: torch.Tensor, y: torch.Tensor, z: TestContainer1):
+        return x + y + z.data_1 + z.data_2[0] + z.data_2[1]
+
+    args = (torch.randn(10, 10),)
+    kwargs = dict(
+        y=torch.randn(10, 10),
+        z=TestContainer1(
+            data_1=torch.randn(10, 10),
+            data_2=(torch.randn(10, 10), torch.randn(10, 10)),
+        ),
+    )
+    flat_inputs = {
+        "args_0": args[0].numpy(),
+        "y": kwargs["y"].numpy(),
+        "z_data_1": kwargs["z"].data_1.numpy(),
+        "z_data_2_0": kwargs["z"].data_2[0].numpy(),
+        "z_data_2_1": kwargs["z"].data_2[1].numpy(),
+    }
+    self._compare_tflite_torch_args_kwargs(SampleModel(), args, kwargs, flat_inputs)
+
+  def test_convert_model_with_args_nested_kwargs_2(self):
+    """
+    Test converting a simple model with both sample_args and nested sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y, z):
+        return x + y + z.data_1 + z.data_2[0][0] + z.data_2[1]
+
+    args = (torch.randn(10, 10),)
+    kwargs = dict(
+        y=torch.randn(10, 10),
+        z=TestContainer1(
+            data_1=torch.randn(10, 10),
+            data_2=[(torch.randn(10, 10),), torch.randn(10, 10)],
+        ),
+    )
+    flat_inputs = {
+        "args_0": args[0].numpy(),
+        "y": kwargs["y"].numpy(),
+        "z_data_1": kwargs["z"].data_1.numpy(),
+        "z_data_2_0_0": kwargs["z"].data_2[0][0].numpy(),
+        "z_data_2_1": kwargs["z"].data_2[1].numpy(),
+    }
+    self._compare_tflite_torch_args_kwargs(SampleModel(), args, kwargs, flat_inputs)
+
+  def test_convert_model_with_args_nested_kwargs_3(self):
+    """
+    Test converting a simple model with both sample_args and nested sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y, z):
+        return x + y + z.data_1 + z.data_2[0]["foo"] + z.data_2[1]
+
+    args = (torch.randn(10, 10),)
+    kwargs = dict(
+        y=torch.randn(10, 10),
+        z=TestContainer1(
+            data_1=torch.randn(10, 10),
+            data_2=(dict(foo=torch.randn(10, 10)), torch.randn(10, 10)),
+        ),
+    )
+    flat_inputs = {
+        "args_0": args[0].numpy(),
+        "y": kwargs["y"].numpy(),
+        "z_data_1": kwargs["z"].data_1.numpy(),
+        "z_data_2_0_foo": kwargs["z"].data_2[0]["foo"].numpy(),
+        "z_data_2_1": kwargs["z"].data_2[1].numpy(),
+    }
+    self._compare_tflite_torch_args_kwargs(SampleModel(), args, kwargs, flat_inputs)
+
+  def _compare_tflite_torch_args_kwargs(self, model, args, kwargs, flat_inputs):
+    model.eval()
+    edge_model = ai_edge_torch.convert(model, args, kwargs)
+    interpreter = tf.lite.Interpreter(model_content=edge_model._tflite_model)
+    runner = interpreter.get_signature_runner("serving_default")
+    input_details = runner.get_input_details()
+    self.assertEqual(input_details.keys(), flat_inputs.keys())
+
+    reference_output = model(*args, **kwargs)
+    tflite_output = edge_model(**flat_inputs)
+    np.testing.assert_almost_equal(reference_output, tflite_output)
+
 
 if __name__ == "__main__":
   unittest.main()

{ai_edge_torch_nightly-0.2.0.dev20240718.dist-info → ai_edge_torch_nightly-0.2.0.dev20240719.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=CNDboRP4zQBpz2hznNCQWcQCARvNXUm3DMa1Dw_XXFg,106
 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=StJHglvx6cii36oi8sj-tZda009e9UqR6ufZOZkP1SY,4137
-ai_edge_torch/convert/conversion_utils.py,sha256=PKXIlSCU-8DhppNBh9ICDNUlEOpV0HgCbt85jDVe3rA,13394
+ai_edge_torch/convert/conversion_utils.py,sha256=TA-fbRApU_wdZYg8VmQSGiH4sm70iITsLRBBi5vODTw,13813
 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
@@ -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=h0vOffr8saDQRkiXljNWDZ17EBjnS4xAtxd8DxETleY,9081
+ai_edge_torch/convert/test/test_convert.py,sha256=itOZDKsh0-0aoly1b1M72M179Yr2BJqtTe6ivueZSc4,12607
 ai_edge_torch/convert/test/test_convert_composites.py,sha256=8UkdPtGkjgSVLCzB_rpM2FmwYuMyt6WE48umX_kr_Sg,7601
 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
@@ -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=kzIulTldq8R9E-lAZsvfSTvLu3FYEX7b9DyYM3qisXM,4485
-ai_edge_torch_nightly-0.2.0.dev20240718.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240718.dist-info/METADATA,sha256=r8YZWPZEhL5gi1oIR9sDZppTSZIuxeHH5isLO4NiSj8,1745
-ai_edge_torch_nightly-0.2.0.dev20240718.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240718.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240718.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240719.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240719.dist-info/METADATA,sha256=X9TaI_3Rxn0rk89P3ZcXJlNtEIUBOhOIIKAncN3Xpos,1745
+ai_edge_torch_nightly-0.2.0.dev20240719.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240719.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240719.dist-info/RECORD,,