ai-edge-torch-nightly 0.2.0.dev20240720__py3-none-any.whl → 0.2.0.dev20240725__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/generative/layers/unet/blocks_2d.py

@@ -125,10 +125,9 @@ class AttentionBlock2D(nn.Module):
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
     if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
       x = input_tensor.view(B, C, H * W)
@@ -181,10 +180,9 @@ class CrossAttentionBlock2D(nn.Module):
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
     if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
       x = input_tensor.view(B, C, H * W)
@@ -222,10 +220,9 @@ class FeedForwardBlock2D(nn.Module):
   def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
     if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
       x = input_tensor.view(B, C, H * W)

{ai_edge_torch_nightly-0.2.0.dev20240720.dist-info → ai_edge_torch_nightly-0.2.0.dev20240725.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240720
+Version: 0.2.0.dev20240725
 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.dev20240720.dist-info → ai_edge_torch_nightly-0.2.0.dev20240725.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
@@ -90,7 +90,7 @@ ai_edge_torch/generative/layers/normalization.py,sha256=M27eW3TcNK20oaXClXtfnu0l
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=12SsCuoRuLNCwnFGe_pHDOZEBwBcqXs87Aj0PaWWw4E,1383
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=dYafGC205QE5CLIbBTCI-7eVvEGZEHzs1toPEhemeDs,3391
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=T70veX57CC9uNidwzoVGzOu-CwzcYMBr1Zk_0bq5UlM,26538
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=4et0TLBtAyVlYxMSJi3-oQoO5npFkOzcCYA927dvm_8,26475
 ai_edge_torch/generative/layers/unet/builder.py,sha256=NmJiZ2-e1wbv9jnvI3VCyUJlONV5ZAOz-RTc7ipAZ5U,1872
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=GU12QEJwO6ukveMR9JRsrhE0YIPKuhk1U81CylmOQTA,9097
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -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.dev20240720.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240720.dist-info/METADATA,sha256=xkXzcnmvTzJRRNOJ2c8JnWS1ZCofdlZiKsW5sa5sDyM,1745
-ai_edge_torch_nightly-0.2.0.dev20240720.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240720.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240720.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/METADATA,sha256=fwnI0u8Q4MJNeahUsCYe34o4YaeoRBJrUieJUk5ogY8,1745
+ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/RECORD,,