ai-edge-torch-nightly 0.3.0.dev20241021__py3-none-any.whl → 0.3.0.dev20241023__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -42,11 +42,18 @@ def lower_by_jax(op, ir_input_names=None):
   return inner
 
 
-_TORCH_XLA2_IMPLS = {
-    key: val.func
-    for key, val in torch_xla2.ops.ops_registry.all_aten_ops.items()
-    if val.is_jax_function
-}
+_TORCH_XLA2_IMPLS = {}
+
+for op, torch_xla2_op in torch_xla2.ops.ops_registry.all_aten_ops.items():
+  if not torch_xla2_op.is_jax_function:
+    continue
+  if isinstance(op, torch._ops.OpOverloadPacket):
+    ops = [getattr(op, overload) for overload in op.overloads()] + [op]
+  else:
+    ops = [op]
+
+  for op in ops:
+    _TORCH_XLA2_IMPLS[op] = torch_xla2_op.func
 
 
 def lower_by_torch_xla2(op):
@@ -229,7 +236,6 @@ lower_by_torch_xla2(torch.ops.aten.transpose_copy)
 lower_by_torch_xla2(torch.ops.aten.triu)
 lower_by_torch_xla2(torch.ops.aten.true_divide)
 lower_by_torch_xla2(torch.ops.aten.trunc)
-lower_by_torch_xla2(torch.ops.aten.unbind)
 lower_by_torch_xla2(torch.ops.aten.unbind_copy)
 lower_by_torch_xla2(torch.ops.aten.unsqueeze)
 lower_by_torch_xla2(torch.ops.aten.unsqueeze.default)
@@ -247,6 +253,11 @@ lower_by_torch_xla2(torch.ops.prims.broadcast_in_dim)
 lower_by_torch_xla2(torch.ops.prims.var)
 
 
+@lower_by_jax(torch.ops.aten.unbind)
+def _aten_copy(self, *args, **kwargs):
+  return _TORCH_XLA2_IMPLS[torch.ops.aten.unbind_copy](self, *args, **kwargs)
+
+
 @lower_by_jax(torch.ops.aten.copy, ir_input_names=["src"])
 def _aten_copy(self, src, **kwargs):
   return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241021.dist-info → ai_edge_torch_nightly-0.3.0.dev20241023.dist-info}/METADATA

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

@@ -3,7 +3,7 @@ ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=SrYveglaiA_DXPoRBqSXClWM1q7853I5ujRorq_MV0M,4251
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=w4N5YyohM9rUsUUc1KQ3ViiuuI09nBTGGwwV0hL9ZcU,706
+ai_edge_torch/version.py,sha256=GjmtIqPIRJ9PpLDTs_m_Jbpvogx7PJ-OK8Bhk-TfbOw,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -168,7 +168,7 @@ ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=dE_qzh-OnCNjWzqs1-PHs5PNlR
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=VvB050UCjB17h6-UNtsaqzVF13MGI01fPFkdmmghTj4,8790
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=v1VdKmL8YLJv3PR9VgyNghO83A25PpTzY2ZUAJqlq3Q,6847
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=RN6BwMHuFj_rFgLCZ6Tu32XHbS2HGjPJeir2nROQ2rA,10517
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=UI_2bbtbkZLBVVcYBXvA6JT1Fg2W_E_OSKYHT-D9xXI,10854
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=1ePJs7oIdUkVdMddFsXMc53qTkEKqGz0ZhQQoNzBa10,2862
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=gqx3n1Mx8pnGQz3nkIF1T_8bkRabXLJBvUoJJn5kOUY,2911
@@ -181,8 +181,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.dev20241021.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/METADATA,sha256=QSx_Fkkgrk5_1B75iuQXcz1p07vsNPenHj59YBQmWuQ,1897
-ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/METADATA,sha256=OfznHDTn9U_zuVoiqlyRXD4CngeLHzKQXKCaXxL8IXs,1897
+ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/RECORD,,