Add direct lowering for aten.abs

PiperOrigin-RevId: 729271527

Author: junjiang-lab

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -215,6 +215,15 @@ def _aten_floor(lctx, x: ir.Value, *, out=None) -> ir.Value:
   return stablehlo.floor(x)
 
 
+# Schema:
+#   - aten::abs(Tensor input) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/main/generated/torch.abs.html
+@lower(torch.ops.aten.abs.default)
+def _aten_abs(lctx, input: ir.Value, *, out=None) -> ir.Value:
+  return stablehlo.abs(input)
+
+
 # Schema:
 #   - aten::cat(Tensor[] tensors, int dim=0) -> Tensor
 # Torch Reference:

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -77,7 +77,6 @@ def lower_by_torch_xla2(op):
 lower_by_torch_xla2(torch.ops.aten._to_copy)
 lower_by_torch_xla2(torch.ops.aten._unsafe_index)
 lower_by_torch_xla2(torch.ops.aten._unsafe_view)
-lower_by_torch_xla2(torch.ops.aten.abs)
 lower_by_torch_xla2(torch.ops.aten.acos)
 lower_by_torch_xla2(torch.ops.aten.acosh)
 lower_by_torch_xla2(torch.ops.aten.add.Scalar)

ai_edge_torch/odml_torch/test/test_core_aten_ops.py

@@ -136,8 +136,11 @@ def _run_export_and_compare(
 
   @parameterized.named_parameters(
       # fmt: off
-      # pyformat: disabledef 
+      # pyformat: disabledef
       ("aten_abs_0", torch.ops.aten.abs, (rnd(torch.float32, (10, 10)),), dict()),
+      ("aten_abs_1", torch.ops.aten.abs, (rnd(torch.float32, (10, 10), -10, 0),), dict()),
+      ("aten_abs_2", torch.ops.aten.abs, (rnd(torch.float32, (10, 10), 0, 10),), dict()),
+      ("aten_abs_3", torch.ops.aten.abs, (rnd(torch.int64, (10, 10), -100, 100),), dict()),
       ("aten_acos_0", torch.ops.aten.acos, (rnd(torch.float32, (10, 10)),), dict()),
       ("aten_acosh_0", torch.ops.aten.acosh, (rnd(torch.float32, (10, 10)),), dict()),
       ("aten_unsqueeze_0", torch.ops.aten.unsqueeze, (rnd(torch.float32, (1, 3, 10)), -2,), dict()),