ai-edge-torch-nightly 0.3.0.dev20241120__py3-none-any.whl → 0.3.0.dev20241121__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -15,13 +15,17 @@
 import math
 from typing import Optional, Union
 
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
 from ai_edge_torch.odml_torch.lowerings import utils
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
 import numpy as np
 import torch
 
-from .registry import lower
+LoweringContext = context.LoweringContext
+lower = registry.lower
 
 
 # add(Tensor self, Tensor other) -> Tensor
@@ -211,6 +215,31 @@ def _aten_floor(lctx, x: ir.Value, *, out=None) -> ir.Value:
   return stablehlo.floor(x)
 
 
+# Schema:
+#   - aten::cat(Tensor[] tensors, int dim=0) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/main/generated/torch.cat.html
+@lower(torch.ops.aten.cat.default)
+def _aten_cat(lctx: LoweringContext, tensors, dim=0):
+  assert tensors
+  non_empty_tensors = [t for t in tensors if np.prod(t.type.shape) != 0]
+  out_meta = lctx.node.meta["tensor_meta"]
+  if not non_empty_tensors:
+    return utils.splat(
+        0,
+        export_utils.torch_dtype_to_ir_element_type(
+            lctx.ir_context, out_meta.dtype
+        ),
+        out_meta.shape,
+    )
+
+  if dim < 0:
+    dim = dim + len(out_meta.shape)
+  dim = ir.IntegerAttr.get(ir.IntegerType.get_signless(64), dim)
+
+  return stablehlo.concatenate(non_empty_tensors, dim)
+
+
 # Schema:
 #   - aten::slice_scatter(Tensor self, Tensor src, int dim=0, SymInt?
 #       start=None, SymInt? end=None, SymInt step=1) -> Tensor

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -105,7 +105,6 @@ lower_by_torch_xla2(torch.ops.aten.bitwise_not)
 lower_by_torch_xla2(torch.ops.aten.bitwise_or)
 lower_by_torch_xla2(torch.ops.aten.bitwise_xor)
 lower_by_torch_xla2(torch.ops.aten.bmm)
-lower_by_torch_xla2(torch.ops.aten.cat)
 lower_by_torch_xla2(torch.ops.aten.ceil)
 lower_by_torch_xla2(torch.ops.aten.clamp.Tensor)
 lower_by_torch_xla2(torch.ops.aten.clamp.default)

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241120.dist-info → ai_edge_torch_nightly-0.3.0.dev20241121.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241120
+Version: 0.3.0.dev20241121
 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.dev20241120.dist-info → ai_edge_torch_nightly-0.3.0.dev20241121.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=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=52sF7t2CBQE8RcB2Hcmo-f6_BLyCW9NzWZ-wTKM9ho4,706
+ai_edge_torch/version.py,sha256=6eLmEn5xqmozokHVWP7j-jjFiQlv2a1aZxDucMzXDh8,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
@@ -178,10 +178,10 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkg
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=drN3L0uTsSjkluKgt6Ngq7b5HLReE_7iAitHpZ9PKqE,5428
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=dE_qzh-OnCNjWzqs1-PHs5PNlRF726qMQKM3tkwAzEs,959
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=z_hPJX9n97d6obcsS9OHXpKqbmw6QqACXgnq5ML6Rhs,9014
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=mxNh20Z4ZeQMu0AAdXnNMXdm2PdAh3RmQPzq2SBpxQs,9954
 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=1JeX3j7Rt3KE7Z2eYRrhtcYgO3EKnRyZFKAUWXw-bsU,10812
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=MY6FFSJKYtD1M1l2q3hDKf3P4NpODqQ4NyWudYe1tTE,10772
 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
@@ -194,8 +194,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.dev20241120.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/METADATA,sha256=1Nv_QeerPRw888sOTf4jHx5Ihu-PJD9rL8GOpRHSTa4,1897
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241120.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241121.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241121.dist-info/METADATA,sha256=AJUg6jkWACMXVy7gopyMvlD0aJfw1BVnkZKbGS9cXX0,1897
+ai_edge_torch_nightly-0.3.0.dev20241121.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
+ai_edge_torch_nightly-0.3.0.dev20241121.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241121.dist-info/RECORD,,