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

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

@@ -41,7 +41,7 @@ _OUTPUT_PATH = flags.DEFINE_string(
 )
 _OUTPUT_NAME_PREFIX = flags.DEFINE_string(
     'output_name_prefix',
-    'deepseek',
+    'amd_llama',
     'The prefix of the output tflite model name.',
 )
 _PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -29,41 +29,46 @@ import torch
 
 _CLIP_CKPT = flags.DEFINE_string(
     'clip_ckpt',
-    None,
+    os.path.join(
+        pathlib.Path.home(),
+        'Downloads/stable-diffusion-v1-5/v1-5-pruned-emaonly.safetensors',
+    ),
     help='Path to source CLIP model checkpoint',
-    required=True,
 )
 
 _DIFFUSION_CKPT = flags.DEFINE_string(
     'diffusion_ckpt',
-    None,
+    os.path.join(
+        pathlib.Path.home(),
+        'Downloads/stable-diffusion-v1-5/v1-5-pruned-emaonly.safetensors',
+    ),
     help='Path to source diffusion model checkpoint',
-    required=True,
 )
 
 _DECODER_CKPT = flags.DEFINE_string(
     'decoder_ckpt',
-    None,
+    os.path.join(
+        pathlib.Path.home(),
+        'Downloads/stable-diffusion-v1-5/v1-5-pruned-emaonly.safetensors',
+    ),
     help='Path to source image decoder model checkpoint',
-    required=True,
 )
 
 _OUTPUT_DIR = flags.DEFINE_string(
     'output_dir',
-    None,
+    '/tmp/sd_tflite',
     help='Path to the converted TF Lite directory.',
-    required=True,
 )
 
 _QUANTIZE = flags.DEFINE_bool(
     'quantize',
     help='Whether to quantize the model during conversion.',
-    default=True,
+    default=False,
 )
 
 _DEVICE_TYPE = flags.DEFINE_string(
     'device_type',
-    'cpu',
+    'gpu',
     help='The device type of the model. Currently supported: cpu, gpu.',
 )
 

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/odml_torch/lowerings/_basic.py

@@ -215,6 +215,15 @@ def _aten_floor(lctx, x: ir.Value, *, out=None) -> ir.Value:
   return stablehlo.floor(x)
 
 
+# Schema:
+#   - aten::abs(Tensor input) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/main/generated/torch.abs.html
+@lower(torch.ops.aten.abs.default)
+def _aten_abs(lctx, input: ir.Value, *, out=None) -> ir.Value:
+  return stablehlo.abs(input)
+
+
 # Schema:
 #   - aten::cat(Tensor[] tensors, int dim=0) -> Tensor
 # Torch Reference:

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -77,7 +77,6 @@ lower_by_torch_xla2(torch.ops.aten._softmax)
 lower_by_torch_xla2(torch.ops.aten._to_copy)
 lower_by_torch_xla2(torch.ops.aten._unsafe_index)
 lower_by_torch_xla2(torch.ops.aten._unsafe_view)
-lower_by_torch_xla2(torch.ops.aten.abs)
 lower_by_torch_xla2(torch.ops.aten.acos)
 lower_by_torch_xla2(torch.ops.aten.acosh)
 lower_by_torch_xla2(torch.ops.aten.add.Scalar)

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.4.0.dev20250220.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.dev20250220
+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.dev20250220.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=ilOEeUufy7PWKexgXSKLjrAG_xeF7RhNEx-8y-4eSYQ,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
@@ -47,7 +47,7 @@ ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/amd_llama_135m/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=urNif89PyCXbdXT5spOeDvdM5luJ-a5HaXHM86v4JnU,2766
-ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=K1cRwdxJWhZ4g97GwI_HwAwU5m5TTIEsjMuGtAlAen8,2563
+ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=JQJrMw50R_9h8tHyhD-GS9WwASBTAnz12tlfVzk9f70,2564
 ai_edge_torch/generative/examples/amd_llama_135m/verify.py,sha256=o13NkFlBgawBsjdJup05VMUjAPvDRAmig6VyEkX8q6U,2426
 ai_edge_torch/generative/examples/deepseek/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=I5eA-XfFdHjYwDsLIjn23T2e-IgnSCQ129-5DOU8j44,2532
@@ -109,9 +109,9 @@ ai_edge_torch/generative/examples/smollm/verify.py,sha256=KpYxVz_lv61YWy6HLfwT68
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 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=GtwKAByEk0ENGEWbUmC2mAAPkbLZ3M5xH1HIToyu8QE,5307
-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/convert_to_tflite.py,sha256=_yk6wVoZm1_FRMFJF5URaPZNNdmMR89fwmKz81BEyao,5601
+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,17 +205,17 @@ 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
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=8mZTp_ybcMO3tDRQdlDP68BVeTw560XsTR4XH-ldTdc,9987
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=Jq8_yAxC7ilzd6tOaRyBsOUEeenFF_EAC5haacZT4Pg,10247
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=VhmeGFnB5hrUsALiVWV96JJOqPDrTIWouHjTvLuT5eU,2477
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=fte81SZxgxeMcI3wWVKSTnUjIxVVilOJ6H3TybXyDmQ,11558
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=h-YHW7xmvt9dpea-7Zj82HW7h5TKzW6GBEE13dIJQ40,11518
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
@@ -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.dev20250220.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250220.dist-info/METADATA,sha256=HZAorXR_re4XSg6gaflqg3YRkdfZcuq7Yxh9K1f1nc4,1966
-ai_edge_torch_nightly-0.4.0.dev20250220.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250220.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250220.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,,