ai-edge-torch-nightly 0.4.0.dev20250221__py3-none-any.whl → 0.4.0.dev20250222__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import ai_edge_torch
 import ai_edge_torch.generative.layers.builder as layers_builder
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 from ai_edge_torch.generative.layers.unet import blocks_2d
@@ -281,7 +282,10 @@ def get_model_config(device_type: str = "cpu") -> unet_cfg.AutoEncoderConfig:
 
   # For now, only turns on StableHLO composite ops on GPU backend for better
   # performance. CPU should also switch to it once the support is done.
-  enable_hlfb = True if device_type == "gpu" else False
+  enable_hlfb = False
+  if device_type == "gpu":
+    ai_edge_torch.config.enable_group_norm_composite = True
+    enable_hlfb = True
 
   norm_config = layers_cfg.NormalizationConfig(
       layers_cfg.NormalizationType.GROUP_NORM,

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import ai_edge_torch
 import ai_edge_torch.generative.layers.builder as layers_builder
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 from ai_edge_torch.generative.layers.unet import blocks_2d
@@ -601,7 +602,10 @@ def get_model_config(
 
   # For now, only turns on StableHLO composite ops on GPU backend for better
   # performance. CPU should also switch to it once the support is done.
-  enable_hlfb = True if device_type == "gpu" else False
+  enable_hlfb = False
+  if device_type == "gpu":
+    ai_edge_torch.config.enable_group_norm_composite = True
+    enable_hlfb = True
 
   # Residual configs.
   residual_norm_config = layers_cfg.NormalizationConfig(

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import torch
 import re
+import torch
 
 
 def _class_fullname(cls):
@@ -35,7 +35,7 @@ def _get_hierarchy(node: torch.fx.Node):
   return hierachy_str
 
 
-def _get_canonical_filename(filename):
+def _get_canonical_filename(filename: str):
   """Remove unnecessary path prefix to make the filename more readable.
 
   This should be factored out so that pattern is a global option that a user
@@ -53,6 +53,29 @@ def _get_canonical_filename(filename):
   return filename
 
 
+def _get_canoical_nodename(node: torch.fx.Node) -> str:
+  """Get the canonical node name from the node's history."""
+
+  history = node.meta.get("from_node", [])
+
+  if len(history) > 1:  # Compatible with torch version under 2.6.0
+    return history[1][0]
+
+  if not hasattr(history[0], "name"):
+    return None
+  names = []
+  while history:
+    names.append(history[0].name)
+    history = history[0].from_node
+
+  # Based on the experiment, the third to last name in the history stack
+  # can be mapped to the original torch node name. The history stack is
+  # generated by tracing the node's transformation history during lowering.
+  if len(names) > 2:
+    return names[-3]
+  return None
+
+
 def build_mlir_file_debuginfo(node: torch.fx.Node):
   """Build the file and line info for the given node's lowerings in MLIR."""
 
@@ -66,16 +89,13 @@ def build_mlir_file_debuginfo(node: torch.fx.Node):
   return _get_canonical_filename(pt_trace.file), int(pt_trace.lineno)
 
 
-def build_nodename_debuginfo(node: torch.fx.Node):
+def build_nodename_debuginfo(node: torch.fx.Node) -> str:
   """Build the fx node name for the given node's lowerings in MLIR."""
-  history = node.meta.get("from_node", [])
-  if not history:
+
+  if not hasattr(node, "meta") or "from_node" not in node.meta:
     return None
-  if len(history) > 1:
-    return history[1][0]
-  if hasattr(history[0], "name"):  # torch 2.6.0+
-    return history[0].name
-  return None
+
+  return _get_canoical_nodename(node)
 
 
 def build_mlir_debuginfo(node: torch.fx.Node):

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250221"
+__version__ = "0.4.0.dev20250222"

{ai_edge_torch_nightly-0.4.0.dev20250221.dist-info → ai_edge_torch_nightly-0.4.0.dev20250222.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=6xZmxpJoHLCndJcctziwhewp4ss--Bi2fe_YyGL2Vag,706
+ai_edge_torch/version.py,sha256=3V323rXqU4LWg1Um9kHwmrGlIkEIet9NlHpxI3IJNyw,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -110,8 +110,8 @@ ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=5M4auM33SgCTODt0VT8TO-EVILruqGDRiNILBPeB83Y,6072
 ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=_yk6wVoZm1_FRMFJF5URaPZNNdmMR89fwmKz81BEyao,5601
-ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=sQKQ-k6H9kG2brgwLsktjCMeN2h0POyfMP6iNsPNKWc,16271
-ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=6W58LxmHHkz2ctgpknQkyoDANZAnE9Byp_svfqLpQf0,34793
+ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=afyHXc86h-ij5zTULmZnM1h313N9VWCyIVriH6pqeSo,16368
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=ylqXOZhYc6XFCaNBKQw0jAnYrCtRFFQKzQzEsFIntvo,34890
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
 ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=GnY3vPZ-obrWuJifuE5bUooKLqAI7v6q71oaTuLKeBE,8778
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=xychak9hdLd6ieXBYEwrK2BkF8NRZWZSSCijIsESpBA,3420
@@ -205,7 +205,7 @@ ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xK
 ai_edge_torch/odml_torch/composite/mark_tensor.py,sha256=U--rwl-XkWKgkdXCXDn6yySug8FR66o1YFUAIoSaWW4,3523
 ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_gLeoXpMcPpV-svXsgN3JbEIjnPVjm0xkpTUdQ,3319
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=3A_lMyj-B-DOhLJG6WmjKvZK5te2rXje8FrfqOhZsN0,959
-ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=JIMCn_fNh5-PgcV5qcklD7aFj0RhNKlvnZ-XQFCOszc,2706
+ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=k9Kas790kpMS5OrVcLzIr48ejAzcc2smrroKAHHM7TQ,3311
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
@@ -230,8 +230,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 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=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.4.0.dev20250221.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250221.dist-info/METADATA,sha256=0V4eezenEyB4Ig9f76FMBy6p1MF3p9Ky--dPmZplYyM,1966
-ai_edge_torch_nightly-0.4.0.dev20250221.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250221.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250221.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250222.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250222.dist-info/METADATA,sha256=nMXtCeRH62h6fhvV2GuzHlG0TDmvPAP_W9KQF6lVsc0,1966
+ai_edge_torch_nightly-0.4.0.dev20250222.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250222.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250222.dist-info/RECORD,,