Add impl for aten.lt.Tensor and tfl.less.

PiperOrigin-RevId: 738084860

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -73,3 +73,8 @@ def _aten_div_tensor_decomp(x, y):
 @register_decomp(torch.ops.aten.gt.Tensor)
 def _aten_gt_tensor_decomp(x, y):
   return torch.ops.tfl.greater(x, y)
+
+
+@register_decomp(torch.ops.aten.lt.Tensor)
+def _aten_lt_tensor_decomp(x, y):
+  return torch.ops.tfl.less(x, y)

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -59,6 +59,11 @@ def tfl_greater(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
   return torch.gt(x, y)
 
 
+@custom_op_with_fake("tfl::less")
+def tfl_less(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+  return torch.lt(x, y)
+
+
 @custom_op_with_fake("tfl::slice")
 def tfl_slice(
     input: torch.Tensor, begin: Sequence[int], size: Sequence[int]

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

@@ -90,6 +90,8 @@ def _assert_export_and_close(
       ("aten_div_Tensor_1", torch.ops.aten.div.Tensor, (rnd(torch.float32, (1, 10)), rnd(torch.float32, (10, 1)),), dict()),
       ("aten_gt_Tensor_0", torch.ops.aten.gt.Tensor, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
       ("aten_gt_Tensor_1", torch.ops.aten.gt.Tensor, (rnd(torch.float32, (1, 10)), rnd(torch.float32, (10, 1)),), dict()),
+      ("aten_lt_Tensor_0", torch.ops.aten.lt.Tensor, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
+      ("aten_lt_Tensor_1", torch.ops.aten.lt.Tensor, (rnd(torch.float32, (1, 10)), rnd(torch.float32, (10, 1)),), dict()),
       # fmt: on
       # pyformat: enable
   )