ai-edge-torch-nightly 0.2.0.dev20240721__py3-none-any.whl → 0.2.0.dev20240726__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/debug/culprit.py

@@ -354,7 +354,7 @@ def _search_model(
     max_granularity: Optional[int] = None,
     enable_fx_minifier_logging: bool = False,
 ) -> Generator[SearchResult, None, None]:
-  """Finds subgraphs in the torch model that satify a certain predicate function provided by the users.
+  """Finds subgraphs in the torch model that satisfy a certain predicate function provided by the users.
 
   Args:
     predicate_f: a predicate function the users specify.
@@ -382,26 +382,33 @@ def _search_model(
   fx_gm, fx_inputs = utils.exported_program_to_fx_graph_module_and_inputs(ep)
   fx_gm = _normalize_getitem_nodes(fx_gm)
 
-  # HACK: temporarily disable XLA_HLO_DEBUG so that fx_minifier won't dump
-  # intermediate stablehlo files to storage.
+  # HACK: temporarily disable XLA_HLO_DEBUG and create_minified_hlo_graph so that
+  # fx_minifier won't dump intermediate stablehlo files to storage.
   # https://github.com/pytorch/pytorch/blob/main/torch/_functorch/fx_minifier.py#L440
   @contextlib.contextmanager
-  def disable_xla_hlo_debug():
+  def disable_minifier_xla_debug():
     xla_hlo_debug_value = None
     if "XLA_HLO_DEBUG" in os.environ:
       xla_hlo_debug_value = os.environ["XLA_HLO_DEBUG"]
       del os.environ["XLA_HLO_DEBUG"]
 
+    create_minified_hlo_graph = torch._functorch.fx_minifier.create_minified_hlo_graph
+    torch._functorch.fx_minifier.create_minified_hlo_graph = (
+        lambda *args, **kwargs: None
+    )
+
     try:
-      yield None
+      yield
     finally:
       if xla_hlo_debug_value is not None:
         os.environ["XLA_HLO_DEBUG"] = xla_hlo_debug_value
 
+      torch._functorch.fx_minifier.create_minified_hlo_graph = create_minified_hlo_graph
+
   found_culprits_num = 0
   while True:
     try:
-      with disable_xla_hlo_debug(), open(os.devnull, "w") as devnull:
+      with disable_minifier_xla_debug(), open(os.devnull, "w") as devnull:
         with contextlib.nullcontext() if enable_fx_minifier_logging else utils.redirect_stdio(
             stdout=devnull,
             stderr=devnull,

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.dev20240721.dist-info → ai_edge_torch_nightly-0.2.0.dev20240726.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240721
+Version: 0.2.0.dev20240726
 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
@@ -27,7 +27,10 @@ Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: safetensors
 Requires-Dist: tabulate
-Requires-Dist: torch >2.3
+Requires-Dist: torch >=2.4.0
+Requires-Dist: torch-xla >=2.4.0
+Requires-Dist: tf-nightly >=2.18.0.dev20240722
+Requires-Dist: ai-edge-quantizer-nightly ==0.0.1.dev20240718
 
 Library that supports converting PyTorch models into a .tflite format, which can
 then be run with TensorFlow Lite and MediaPipe.  This enables applications for

{ai_edge_torch_nightly-0.2.0.dev20240721.dist-info → ai_edge_torch_nightly-0.2.0.dev20240726.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
@@ -27,7 +27,7 @@ ai_edge_torch/convert/test/test_convert_composites.py,sha256=8UkdPtGkjgSVLCzB_rp
 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
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
-ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
+ai_edge_torch/debug/culprit.py,sha256=lN2N_J3EJTPcEH42xK6kvs4yOqZtyvvecSc4aidlog4,14556
 ai_edge_torch/debug/utils.py,sha256=hjVmQVVl1dKxEF0D6KB4a3ouQ3wBkTsebOX2YsUObZM,1430
 ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/debug/test/test_culprit.py,sha256=9An_n9p_RWTAYdHYTCO-__EJlbnjclCDo8tDhOzMlwk,3731
@@ -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.dev20240721.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240721.dist-info/METADATA,sha256=TJYFNAxXQkRwt9I_0OqpUOS3opWBU5i-ioMwsicD7cY,1745
-ai_edge_torch_nightly-0.2.0.dev20240721.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240721.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240721.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/METADATA,sha256=QstgV90JdvfDL5-OAb_Z4dzWItNf6qT-gNfDtzdoTok,1889
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/RECORD,,