Support torch.ops.aten.sym_size.int and torch Tensor.reshape.

PiperOrigin-RevId: 750232210

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -269,15 +269,41 @@ def _tfl_transpose_lowering(
 def _tfl_reshape_lowering(
     lctx: LoweringContext,
     x: ir.Value,
-    shape: Sequence[int],
+    shape: Sequence[int | ir.Value],
 ) -> ir.Value:
-  constant_shape = lowering_utils.numpy_array_constant(
-      np.array(shape, dtype=np.int32)
-  )
+  # Check if all elements in the shape sequence are integers.
+  if not shape or all(isinstance(dim, int) for dim in shape):
+    # If all are integers, create a constant numpy array.
+    # Assuming int32 is the required type for TFLite shape tensors.
+    shape_ir_value = lowering_utils.numpy_array_constant(
+        np.array(shape, dtype=np.int32)
+    )
+  else:
+    # Handle mixed int and ir.Value shape sequence
+    processed_dims = []
+    for dim in shape:
+      if isinstance(dim, int):
+        # Convert int to a constant 1D tensor
+        shape_ir_value = lowering_utils.numpy_array_constant(
+            np.array([dim], dtype=np.int32)
+        )
+        processed_dims.append(shape_ir_value)
+      else:
+        assert isinstance(dim, ir.Value)
+        # Convert ir.Value to a constant 1D tensor
+        new_type = ir.RankedTensorType.get([1], dim.type.element_type)
+        reshape_dim = stablehlo.reshape(new_type, dim)
+        processed_dims.append(reshape_dim)
+
+    shape_ir_value = stablehlo.concatenate(
+        processed_dims,
+        dimension=0,
+    )
+
   return _ir_operation(
       "tfl.reshape",
       results=lowering_utils.node_meta_to_ir_types(lctx.node),
-      operands=[x, constant_shape],
+      operands=[x, shape_ir_value],
   )
 
 

ai_edge_torch/odml_torch/experimental/torch_tfl/test/test_torch_tfl_impls.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Numerical validation tests for torch ops and Torch-TFL ops."""
 
+from typing import Any, Dict, Sequence
+
 import ai_edge_torch
 from ai_edge_torch import testing
 from ai_edge_torch.odml_torch.experimental import torch_tfl
@@ -53,14 +55,23 @@ def setUp(self):
     torch.manual_seed(0)
 
   def _assert_export_and_close(
-      self, func, args, kwargs, atol=1e-3, rtol=1e-5, equal_nan=True
+      self,
+      func,
+      args,
+      kwargs,
+      dynamic_shapes: Dict[str, Any] | Sequence[Any] | None = None,
+      atol=1e-3,
+      rtol=1e-5,
+      equal_nan=True,
   ):
     """Assert func, args, and kwargs can be lowered and pass numerical validation."""
     with self.subTest("torch_eval"):
       expected = func(*args, **kwargs)
 
       with self.subTest("export_and_decompse"):
-        exported_program = export_with_tensor_inputs_only(func, args, kwargs)
+        exported_program = export_with_tensor_inputs_only(
+            func, args, kwargs, dynamic_shapes
+        )
         exported_program = exported_program.run_decompositions(
             torch_tfl.decomps
         )
@@ -81,7 +92,9 @@ def _assert_export_and_close(
 
         with self.subTest("convert_eval"):
           args, kwargs = exported_program.example_inputs
-          edge_model = ai_edge_torch.convert(exported_program.module(), args)
+          edge_model = ai_edge_torch.convert(
+              exported_program.module(), args, dynamic_shapes=dynamic_shapes
+          )
           actual = edge_model(*args, **kwargs)
 
           with self.subTest("torch_convert_eval_diff:" + str(atol)):
@@ -142,9 +155,41 @@ def _assert_export_and_close(
       # fmt: on
       # pyformat: enable
   )
-  def test_op(self, op, args, kwargs):
+  def test_op(
+      self,
+      op,
+      args,
+      kwargs,
+  ):
     self._assert_export_and_close(op, args, kwargs)
 
+  @parameterized.named_parameters(
+      # fmt: off
+      # pyformat: disabledef
+      ("reshape_without_dynamic_shape_0", (rnd(torch.float32, (10, 2, 3)),), dict(), None),
+      ("reshape_with_dynamic_shape_1", (rnd(torch.float32, (10, 2, 3)),), dict(), ((torch.export.Dim("batch"), None, None),)),
+      ("reshape_with_dynamic_shape_2", (rnd(torch.float32, (10, 2, 3)),), dict(), ({0: torch.export.Dim("batch")},)),
+      ("reshape_with_dynamic_shape_3", (rnd(torch.float32, (10, 2, 3)),), dict(), ((torch.export.Dim("batch"), torch.export.Dim("height"), torch.export.Dim("width")),)),
+      ("reshape_with_dynamic_shape_4", (rnd(torch.float32, (10, 2, 3)),), dict(), ({0: torch.export.Dim("batch"), 1: torch.export.Dim("height"), 2: torch.export.Dim("width")},)),
+      # fmt: on
+      # pyformat: enable
+  )
+  def test_reshape_op(
+      self,
+      args,
+      kwargs,
+      dynamic_shapes: Dict[str, Any] | Sequence[Any] | None = None,
+  ):
+
+    class ReshapeModel(torch.nn.Module):
+
+      def forward(self, x):
+        x = x + x
+        x = x.reshape([x.shape[0], x.shape[1] * x.shape[2]])
+        return x
+
+    self._assert_export_and_close(ReshapeModel(), args, kwargs, dynamic_shapes)
+
 
 if __name__ == "__main__":
   googletest.main()

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 @@
 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 @@ def lower_by_torch_xla2(op):
 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 forward(self, *export_args):
 
   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,
   )