ai-edge-torch-nightly 0.5.0.dev20250422__py3-none-any.whl → 0.5.0.dev20250423__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -37,6 +37,7 @@ def _run_convert_passes(
   passes = [
       fx_passes.CastInputsBf16ToF32Pass(),
       fx_passes.BuildInterpolateCompositePass(),
+      fx_passes.CanonicalizePass(),
       fx_passes.OptimizeLayoutTransposesPass(),
       fx_passes.CanonicalizePass(),
       fx_passes.BuildAtenCompositePass(),

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import math
+import operator
 from typing import Optional, Union
 
 from ai_edge_torch.odml_torch import export_utils
@@ -24,6 +25,7 @@ from jax._src.lib.mlir.dialects import hlo as stablehlo
 import numpy as np
 import torch
 
+
 LoweringContext = context.LoweringContext
 lower = registry.lower
 
@@ -320,3 +322,22 @@ def _aten_to_copy(
       ),
       x,
   )
+
+
+# Schema:
+#   - aten::sym_size.int(Tensor self, int dim) -> SymInt
+@lower(torch.ops.aten.sym_size.int)
+def _aten_sym_size_int(lctx, x: ir.Value, dim: int):
+  return stablehlo.get_dimension_size(x, dim)
+
+
+# Lowering for the multiplication operator (`*`).
+# Handles cases where one operand is an integer (scalar) and the other is a
+# tensor, broadcasting the scalar to the tensor's shape before multiplication.
+@lower(operator.mul)
+def _operator_mul(lctx, self: int | ir.Value, other: int | ir.Value):
+  if isinstance(self, int) and isinstance(other, ir.Value):
+    self = utils.splat(self, other.type.element_type, other.type.shape)
+  if isinstance(other, int) and isinstance(self, ir.Value):
+    other = utils.splat(other, self.type.element_type, self.type.shape)
+  return stablehlo.multiply(self, other)

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -218,7 +218,6 @@ lower_by_torch_xla2(torch.ops.aten.stack)
 lower_by_torch_xla2(torch.ops.aten.sub.Scalar)
 lower_by_torch_xla2(torch.ops.aten.sub.Tensor)
 lower_by_torch_xla2(torch.ops.aten.sum)
-lower_by_torch_xla2(torch.ops.aten.sym_size)
 lower_by_torch_xla2(torch.ops.aten.t)
 lower_by_torch_xla2(torch.ops.aten.tan)
 lower_by_torch_xla2(torch.ops.aten.tanh)

ai_edge_torch/testing/export.py

@@ -15,7 +15,7 @@
 """Torch export utilities for testing."""
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, Dict, Sequence
 
 import torch
 from torch.utils import _pytree as pytree
@@ -25,6 +25,7 @@ def export_with_tensor_inputs_only(
     model: Callable[..., Any],
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
+    dynamic_shapes: Dict[str, Any] | Sequence[Any] | None = None,
 ) -> torch.export.ExportedProgram:
   """Exports a PyTorch model, treating only tensor inputs as export inputs.
 
@@ -76,8 +77,12 @@ def export_with_tensor_inputs_only(
 
   export_args = tuple(export_args)
   export_kwargs = {}
+  # Need to wrap dynamic_shapes in a tuple to match the inputs structure of
+  # ModuleWrapper.
+  dynamic_shapes = (dynamic_shapes,) if dynamic_shapes else None
   return torch.export.export(
       ModuleWrapper(model, args, kwargs).eval(),
       export_args,
       export_kwargs,
+      dynamic_shapes=dynamic_shapes,
   )

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250422"
+__version__ = "0.5.0.dev20250423"

{ai_edge_torch_nightly-0.5.0.dev20250422.dist-info → ai_edge_torch_nightly-0.5.0.dev20250423.dist-info}/METADATA

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

@@ -2,9 +2,9 @@ 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=8A53NhGxkOePrHhrJE-p0LgdIpk8vT_45m6C94N8uEw,706
+ai_edge_torch/version.py,sha256=DjzQwP8czvLmUu-dJhnWVQJHOuaOqJJKuH2_TOViMvg,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=0gpwEjlTue5RttDerzM5SVOUnY8g16444yL2YIFBx-E,5485
+ai_edge_torch/_convert/conversion.py,sha256=dOr3TUfF0UCvkmlUrMqKvgaN4jh3lJ9XFuO-sHaAmIw,5521
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
@@ -225,11 +225,11 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNi
 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=4syWstepGiw3IKa8O7lciXywY7RFJ7OCWFMU1Lg3h-s,10777
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=fEWjIdEpDIqT1EYLZE13O9A41OuaNdbfBrv3vNxS9gI,11601
 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=JRGLXW8EQ1L-vdiVTkD1kb4AnTU05eRwZ7Ke010hZmg,11473
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=tkaDo232HjuZvJHyua0n6tdHecifUuVzclJAGq7PPYs,11428
 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
@@ -242,11 +242,11 @@ ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV
 ai_edge_torch/quantize/pt2e_quantizer_utils.py,sha256=nuO3w9gOj9sKcsTBBexVDw8UZnd06KsjNrFr_gyNaiA,36710
 ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9PphCRdO8o,3172
 ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPGhPbo,833
-ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
+ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.5.0.dev20250422.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250422.dist-info/METADATA,sha256=9iAyhGknDER60qxgpUn6ZSR6Nr7LALvvHjgOf0WnYtg,2051
-ai_edge_torch_nightly-0.5.0.dev20250422.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250422.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250422.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250423.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250423.dist-info/METADATA,sha256=PGzcX4WVfFW0wE0TSKLAuRB94iemrNff4L8CL_VUMnQ,2051
+ai_edge_torch_nightly-0.5.0.dev20250423.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250423.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250423.dist-info/RECORD,,