ai-edge-torch-nightly 0.2.0.dev20240723__py3-none-any.whl → 0.2.0.dev20240724__py3-none-any.whl

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -113,6 +113,10 @@ def is_4d(node: Node):
   val = node.meta.get("val")
   if val is None:
     return False
+
+  if isinstance(val, (list, tuple)) and val:
+    val = val[0]
+
   if not hasattr(val, "shape"):
     return False
 
@@ -168,7 +172,6 @@ def _all_layout_sensitive_inputs_are_4d_checker(node: Node):
 
 
 @nhwcable_node_checkers.register(aten._native_batch_norm_legit_no_training)
-@nhwcable_node_checkers.register(aten.native_group_norm)
 def _aten_norm_checker(node):
   val = node.meta.get("val")
   if not isinstance(val, (list, tuple)) or not val or not hasattr(val[0], "shape"):
@@ -176,6 +179,18 @@ def _aten_norm_checker(node):
   return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
 
 
+@nhwcable_node_checkers.register(aten.native_group_norm)
+def _aten_native_group_norm_checker(node):
+  val = node.meta.get("val")
+  if not isinstance(val, (list, tuple)) or not val or not hasattr(val[0], "shape"):
+    return NHWCable(can_be=False, must_be=False)
+  if len(node.args) >= 3 and (node.args[1] is not None or node.args[2] is not None):
+    # Disable NHWC rewriter due to precision issue with weight and bias.
+    # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
+    return NHWCable(can_be=False, must_be=False)
+  return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
+
+
 # ==== Ops must be NCHW
 
 

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import operator
 import os
 from typing import Optional, Tuple, Union
 
@@ -274,6 +275,14 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
         graph_module = layout_partitioners.greedy.partition(graph_module)
 
     graph = graph_module.graph
+    for node in list(graph.nodes):
+      if node.target == operator.getitem:
+        # force the layout mark of a getitem node to follow its producer.
+        if layout_mark.is_nchw_node(node.args[0]):
+          layout_mark.mark_as_nchw_node(node)
+        else:
+          layout_mark.mark_as_nhwc_node(node)
+
     for node in list(graph.nodes):
       if layout_mark.is_nhwc_node(node):
         for input_node in layout_check.get_layout_sensitive_inputs(node):

{ai_edge_torch_nightly-0.2.0.dev20240723.dist-info → ai_edge_torch_nightly-0.2.0.dev20240724.dist-info}/METADATA

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

@@ -12,11 +12,11 @@ ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=6m_vc
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=UX6dJsxCqSkftXXvNBV-i7Bjk6H7qTyqzUnE640Itfg,1673
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=VA9bekxPVhLk4MYlIRXnOzrSnbCtUmGj7OQ_fJcKQtc,795
-ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=pG-zLvO5vGs3gjNXa3RxGNwvC-_Azei2anxe2VdKsnY,6870
+ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=kV4Md7YBQIS-A8Dp4O-8SugNHAfDjVIBHvnldVPpHV0,7483
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=uXCHC23pWN-3JmDtAErWbSUnL8jjlQgUAy4gqtfDsQU,1560
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=RAgU31B98PQmXEIM3GOjgS0q9aRe2whJhGXpW2EjoqY,12438
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=tCx7J-WIFnxFCeRBtqJ159jWLgK9_9DCJrR4mkeBuYE,982
-ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=cfY6RTWQTGXNoQxKHaDcBYR9QdkVQXOWjKhuxvglocw,10383
+ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=7SuNFBfQWt42tEcUUxpQVnWhy-ByMuLHi2VtH8Kc1g4,10708
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/utils.py,sha256=bItkXVaPA9THcFypAmqldpkLuD8WpOFmKlhVbBJJkPk,2076
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py,sha256=B-zisphkH7aRCUOJNdwHnTA0fQXuDpN08q3Qjy5bL6E,715
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=8uHJbIwPMTgeSfYVba163pkXSQkHLxFwar_8A1AhgAM,2279
@@ -125,8 +125,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.dev20240723.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240723.dist-info/METADATA,sha256=X_OA68Nyy0x4kv5FoE92mxTAegD55SCIMmJ2fmPhcoQ,1745
-ai_edge_torch_nightly-0.2.0.dev20240723.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240723.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240723.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240724.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240724.dist-info/METADATA,sha256=pdYXzmVJ02-GQnxWNMHtYvAq4LaTOpfFREjbRcpNSmI,1745
+ai_edge_torch_nightly-0.2.0.dev20240724.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240724.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240724.dist-info/RECORD,,