Add impl for aten.arange.start_step and lowering.

PiperOrigin-RevId: 766755082

Author: junjiang-lab

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py

@@ -79,7 +79,9 @@ def edges(self):
 
   @property
   def graph(self):
-    edges = np.array(self.edges)
+    # Ensure edges is a 2D array with shape (N, 3) and int32 dtype.
+    # If self.edges is empty, this will result in an array with shape (0, 3).
+    edges = np.array(self.edges, dtype=np.int32).reshape(-1, 3)
     return scipy.sparse.csr_matrix(
         (
             np.minimum(edges[:, 2], MinCutSolver.INF_COST),

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -160,6 +160,13 @@ def _aten_view_decomp(x, shape: Sequence[int]):
   return torch.ops.tfl.reshape(x, shape)
 
 
+@register_decomp(torch.ops.aten.arange.start_step)
+def _aten_arange_start_step_decomp(
+    start, end, step=1, dtype=None, layout=None, device=None, pin_memory=None
+):
+  return torch.ops.tfl.range(start, end, step)
+
+
 @register_decomp(torch.ops.aten._softmax.default)
 def _aten__softmax_decomp(
     x, dim: int, half_to_float: bool  # pylint: disable=unused-argument

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -339,6 +339,75 @@ def _tfl_reshape_lowering(
   )
 
 
+@lower(torch.ops.tfl.range.default)
+def _tfl_range_lowering(
+    lctx: LoweringContext,
+    start: int | float | ir.Value,
+    limit: int | float | ir.Value,
+    delta: int | float | ir.Value = 1,
+) -> ir.Value:
+  tensor_meta = lctx.node.meta.get("tensor_meta") or lctx.node.meta.get("val")
+  output_torch_dtype = tensor_meta.dtype
+
+  original_mlir_output_types = lowering_utils.node_meta_to_ir_types(lctx.node)
+  if not original_mlir_output_types or not isinstance(
+      original_mlir_output_types[0], ir.RankedTensorType
+  ):
+    raise ValueError(
+        "tfl.range output type is not a RankedTensorType as expected."
+    )
+
+  original_mlir_output_type = original_mlir_output_types[0]
+  original_output_shape = original_mlir_output_type.shape
+  original_output_element_type = original_mlir_output_type.element_type
+  tflite_op_internal_element_type = (
+      lowering_utils.torch_dtype_to_ir_element_type(output_torch_dtype)
+  )
+
+  operands = []
+  for val_py_scalar in [
+      start,
+      limit,
+      delta,
+  ]:
+    if isinstance(val_py_scalar, ir.Value):
+      operands.append(val_py_scalar)
+    else:
+      numpy_scalar_0d = (
+          torch.tensor(val_py_scalar, dtype=output_torch_dtype)
+          .detach()
+          .numpy()
+      )
+      scalar_tensor_val = lowering_utils.numpy_array_constant(numpy_scalar_0d)
+      operands.append(scalar_tensor_val)
+
+  # Define the result type that the tfl.range *kernel* (the custom op) will
+  # produce.
+  tfl_op_kernel_output_type = ir.RankedTensorType.get(
+      original_output_shape, tflite_op_internal_element_type
+  )
+
+  tfl_range_op_val = _ir_operation(
+      "tfl.range",
+      results=[tfl_op_kernel_output_type],
+      operands=operands,
+  )
+
+  # The _tfl_range_lowering function must return a value of the
+  # original_mlir_output_type.
+  # If the tfl.range op's internal element type is different from the
+  # original_output_element_type, we need to convert.
+  if tflite_op_internal_element_type != original_output_element_type:
+    # Convert the tfl.range output to the original expected type.
+    final_output_val = stablehlo.convert(
+        original_mlir_output_type, tfl_range_op_val
+    )
+  else:
+    final_output_val = tfl_range_op_val
+
+  return final_output_val
+
+
 @lower(torch.ops.tfl.softmax.default)
 def _tfl_softmax_lowering(
     lctx: LoweringContext,

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -189,6 +189,15 @@ def tfl_reshape_fake(input: torch.Tensor, shape: Sequence[int]) -> torch.Tensor:
   return torch.empty(inferred_shape, dtype=input.dtype)
 
 
+@custom_op_with_fake(
+    "tfl::range", schema="(Scalar start, Scalar limit, Scalar delta) -> Tensor"
+)
+def tfl_range(
+    start: int | float, limit: int | float, delta: int | float
+) -> torch.Tensor:
+  return torch.arange(start, limit, delta)
+
+
 @custom_op_with_fake("tfl::softmax")
 def tfl_softmax(x: torch.Tensor) -> torch.Tensor:
   return torch.nn.functional.softmax(x, dim=-1)

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

@@ -154,6 +154,10 @@ def _assert_export_and_close(
       ("aten_gelu_3", torch.ops.aten.gelu.default, (rnd(torch.float32, (1, 10)),), dict(approximate="tanh")),
       ("aten_permute_0", torch.ops.aten.permute.default, (rnd(torch.float32, (10, 10)), [0, 1],), dict()),
       ("aten_permute_1", torch.ops.aten.permute.default, (rnd(torch.float32, (1, 10)), [0, 1],), dict()),
+      ("aten_arange_start_step_0", torch.ops.aten.arange.start_step, (0, 100, 5), dict()),
+      ("aten_arange_start_step_1", torch.ops.aten.arange.start_step, (0, 100), dict()),
+      ("aten_arange_start_step_2", torch.ops.aten.arange.start_step, (100, 0, -1), dict()),
+      ("aten_arange_start_step_3", torch.ops.aten.arange.start_step, (0, 10.5, 0.5), dict()),
       ("aten_view_0", torch.ops.aten.view.default, (rnd(torch.float32, (10, 10)), [1, 100],), dict()),
       ("aten_view_1", torch.ops.aten.view.default, (rnd(torch.float32, (1, 10)), [10, 1],), dict()),
       ("aten_view_2", torch.ops.aten.view.default, (rnd(torch.float32, (10, 10)), [2, 5, 10],), dict()),