ai-edge-torch-nightly 0.3.0.dev20250211__py3-none-any.whl → 0.3.0.dev20250212__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -56,11 +56,6 @@ _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
-_PIXEL_VALUES_SIZE = flags.DEFINE_multi_integer(
-    'pixel_values_size',
-    [3, 224, 224],
-    'The size of prefill pixel values except the batch dimension.',
-)
 _QUANTIZE = flags.DEFINE_bool(
     'quantize',
     True,
@@ -75,12 +70,15 @@ def main(_):
       kv_cache_max_len=_KV_CACHE_MAX_LEN.value,
   )
 
+  config = pytorch_model.image_encoder.config.image_embedding
   converter.convert_to_tflite(
       pytorch_model,
       output_path=_OUTPUT_PATH.value,
       output_name_prefix=f'{_OUTPUT_NAME_PREFIX.value}_{_VERSION.value}',
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
-      pixel_values_size=torch.Size(_PIXEL_VALUES_SIZE.value),
+      pixel_values_size=torch.Size(
+          [1, config.channels, config.image_size, config.image_size]
+      ),
       quantize=_QUANTIZE.value,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -136,9 +136,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      # TODO: b/377051577 - Once RemoveSDPACompositeZeroMaskPass is removed,
-      # enable_hlfb can be set to True. See b/383865404#comment3 for details.
-      # enable_hlfb=True,
+      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/utilities/converter.py

@@ -145,7 +145,7 @@ def _export_helper(
     prefill_input_pos_list.append(torch.arange(0, seq_len, dtype=torch.int))
 
   prefill_pixel_values = (
-      torch.full((1,) + pixel_values_size, 0, dtype=torch.float32)
+      torch.full(pixel_values_size, 0, dtype=torch.float32)
       if pixel_values_size
       else None
   )

ai_edge_torch/odml_torch/debuginfo/__init__.py

@@ -12,5 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ._build import build_mlir_debuginfo, build_mlir_file_debuginfo
+"""Debug info generation for ODML Torch."""
+
+from . import _build
 from ._op_polyfill import write_mlir_debuginfo_op
+
+build_nodename_debuginfo = _build.build_nodename_debuginfo
+build_mlir_file_debuginfo = _build.build_mlir_file_debuginfo
+build_mlir_debuginfo = _build.build_mlir_debuginfo

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -40,10 +40,16 @@ def _get_canonical_filename(filename):
 
   This should be factored out so that pattern is a global option that a user
   can override.
+
+  Args:
+    filename: The filename to canonicalize.
+
+  Returns:
+    The canonicalized filename.
   """
 
-  # TODO: We should add a config option to provide a regex to strip from the
-  # debug info. Currently absolute path is used.
+  # TODO(yijieyang): We should add a config option to provide a regex to strip
+  # from the debug info. Currently absolute path is used.
   return filename
 
 
@@ -55,9 +61,23 @@ def build_mlir_file_debuginfo(node: torch.fx.Node):
 
   # Note: This uses internal APIs and may break in the future.
   pt_trace = torch.fx.graph._parse_stack_trace(node.stack_trace)
+  if pt_trace is None:
+    return None, None
   return _get_canonical_filename(pt_trace.file), int(pt_trace.lineno)
 
 
+def build_nodename_debuginfo(node: torch.fx.Node):
+  """Build the fx node name for the given node's lowerings in MLIR."""
+  history = node.meta.get("from_node", [])
+  if not history:
+    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
+
+
 def build_mlir_debuginfo(node: torch.fx.Node):
   """Build the debuginfo string for the given node's lowerings in MLIR."""
 

ai_edge_torch/odml_torch/export.py

@@ -88,6 +88,21 @@ class LoweringInterpreter(torch.fx.Interpreter):
     self.outputs = None
 
   def _build_loc(self, node: torch.fx.Node):
+    """Build MLIR location for the given node.
+
+    The location contains:
+    - layer info
+    - fx node name
+    - file and line info
+
+    Currently it's still under development and format is subject to change.
+
+    Args:
+      node: The torch.fx.Node to build the location for.
+
+    Returns:
+      The MLIR location for the given node.
+    """
 
     info = debuginfo.build_mlir_debuginfo(node)
     if info is None:
@@ -98,7 +113,12 @@ class LoweringInterpreter(torch.fx.Interpreter):
     if file is not None:
       fileinfo = ir.Location.file(filename=file, line=line, col=0)
 
-    return ir.Location.name(name=info, childLoc=fileinfo)
+    node_name = debuginfo.build_nodename_debuginfo(node)
+    nodeinfo = None
+    if node_name is not None:
+      nodeinfo = ir.Location.name(name=node_name, childLoc=fileinfo)
+
+    return ir.Location.name(name=info, childLoc=nodeinfo)
 
   def run_node(self, node: torch.fx.Node):
     loc = self._build_loc(node)

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20250211"
+__version__ = "0.3.0.dev20250212"

{ai_edge_torch_nightly-0.3.0.dev20250211.dist-info → ai_edge_torch_nightly-0.3.0.dev20250212.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250211
+Version: 0.3.0.dev20250212
 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.3.0.dev20250211.dist-info → ai_edge_torch_nightly-0.3.0.dev20250212.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=SqHTGIWZweve6ni4shLmLl85BwAAptNTCG646yd6Hk8,706
+ai_edge_torch/version.py,sha256=gYCO29a18zxFgm6z4IHCArhNE29IvlL_sxhEL6O4ECw,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
@@ -73,10 +73,10 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=pyxRGgMxrn
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3TifYv-SLOSheaWKv2AH0iKDrc,2441
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=a6ISb96xhEJc1TtaFGCUiA4msKedPTAeMvkWrfIklx4,2792
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=APQymtr3n2k6-e8wvn3kVrli0qiElduYIkHeahcoSA0,2743
 ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=z658dW_D0Iqvo6xnh4vG7_o17-Fufndyis8Rq5yafJY,5439
 ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=GZa0Ou_DvOijB2nTL_jRvGbn0_dvJPosQAPf47yqicw,5988
-ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=7K1xl64UvoHaYmqWjIbahwXHfppwTQ8sN7JrpGKX1XQ,5771
+ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
 ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=x1mgRtVLxkCTvlkPow3y7ADoGTjUh5uc5pF46mxatLw,6099
 ai_edge_torch/generative/examples/paligemma/verify.py,sha256=HLcu1fWMtFFFONAqVW94rOBqq4XvFHtatX3JFGOsfZw,5345
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
@@ -172,7 +172,7 @@ ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728Fc
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/bmm_4d.py,sha256=2BMOYiFVUsl-bjxmLkrX4N7kpO0CnhB7eDYxm_iBCr8,2533
-ai_edge_torch/generative/utilities/converter.py,sha256=6siSpCvH_cLV-eP40lkF_AqjBpYv68xeMRQ722fKgE0,8065
+ai_edge_torch/generative/utilities/converter.py,sha256=K9taR0KY59dvfU_jO1yBe_p7w8lDns1Q3U6oJTTKZzM,8058
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
 ai_edge_torch/generative/utilities/model_builder.py,sha256=5WqcxpeTdt51nVoUwt9g5kKB5wQKj2eYbiaz7k6Ofxg,6815
@@ -197,14 +197,14 @@ ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=LDyZUehM1lmT3y2bGeA94rMGRUTLxzIUm4DTlCA8tQc,13640
+ai_edge_torch/odml_torch/export.py,sha256=7l8R0DEq_vfns8iWpruMlIyaIKZAFzoAy369-7iRrl0,14164
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
 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=DoE3HgAtV_GNKGBDGzH2Lb7JUHvyH7TUqWbDZIObr34,789
-ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=sjpYeqgdbDmD7lhp80yc8jfWq-HxX3xuQ58ND8ZeU-I,2213
+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/_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
@@ -229,8 +229,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.3.0.dev20250211.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250211.dist-info/METADATA,sha256=wT5v4PcnaE4IJmA-3d1W-0U91ETCBsWZyaT8VWiVI-c,1966
-ai_edge_torch_nightly-0.3.0.dev20250211.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250211.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250211.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250212.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250212.dist-info/METADATA,sha256=uCdhC0TvTD1j2cNaIsI1l_uB_FVHMml_jDr2vUzzKHw,1966
+ai_edge_torch_nightly-0.3.0.dev20250212.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250212.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250212.dist-info/RECORD,,