Add impl for aten._softmax.default, tfl.softmax and lowering.

PiperOrigin-RevId: 743639346

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -114,3 +114,20 @@ def _aten_gelu_decomp(x, approximate="none"):
 @register_decomp(torch.ops.aten.permute.default)
 def _aten_permute_decomp(x, dims: Sequence[int]):
   return torch.ops.tfl.transpose(x, dims)
+
+
+@register_decomp(torch.ops.aten._softmax.default)
+def _aten__softmax_decomp(
+    x, dim: int, half_to_float: bool  # pylint: disable=unused-argument
+):
+  if dim == -1 or dim == x.dim() - 1:
+    return torch.ops.tfl.softmax(x)
+  else:
+    dims = list(range(x.dim()))
+    # Transpose the input by swapping the dim with the last dimension.
+    dims[dim], dims[-1] = dims[-1], dims[dim]
+    x_permuted = torch.ops.tfl.transpose(x, dims)
+    # Compute the softmax on the last dimension.
+    softmax_result = torch.ops.tfl.softmax(x_permuted)
+    # Transpose the result back to the original dimensions.
+    return torch.ops.tfl.transpose(softmax_result, dims)

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -263,3 +263,19 @@ def _tfl_transpose_lowering(
       results=lowering_utils.node_meta_to_ir_types(lctx.node),
       operands=[x, constant_perm],
   )
+
+
+@lower(torch.ops.tfl.softmax.default)
+def _tfl_softmax_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    beta: float = 1.0,
+) -> ir.Value:
+  return _ir_operation(
+      "tfl.softmax",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x],
+      attributes={
+          "beta": ir.FloatAttr.get(ir.F32Type.get(), beta),
+      },
+  )

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -102,6 +102,11 @@ def tfl_transpose(input: torch.Tensor, perm: Sequence[int]) -> torch.Tensor:
   return torch.permute(input, perm).clone()
 
 
+@custom_op_with_fake("tfl::softmax")
+def tfl_softmax(x: torch.Tensor) -> torch.Tensor:
+  return torch.nn.functional.softmax(x)
+
+
 @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

@@ -130,6 +130,12 @@ 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__softmax_0", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), -1, False), dict()),
+      ("aten__softmax_1", torch.ops.aten._softmax.default, (rnd(torch.float32, (1, 10)), -1, False), dict()),
+      ("aten__softmax_2", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), 0, False), dict()),
+      ("aten__softmax_3", torch.ops.aten._softmax.default, (rnd(torch.float32, (1, 10)), 0, False), dict()),
+      ("aten__softmax_4", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), 1, False), dict()),
+      ("aten__softmax_5", torch.ops.aten._softmax.default, (rnd(torch.float32, (1, 10)), 1, False), dict()),
       # fmt: on
       # pyformat: enable
   )