Add basic softmax implementation that handles limited axis dim length

Author: Ziminli

src/ntops/__init__.py

@@ -7,5 +7,6 @@
 from ntops.gelu import gelu
 from ntops.mm import mm
 from ntops.mul import mul
+from ntops.softmax import softmax
 
-__all__ = ["abs", "add", "addmm", "bmm", "div", "exp", "gelu", "mm", "mul"]
+__all__ = ["abs", "add", "addmm", "bmm", "div", "exp", "gelu", "mm", "mul", "softmax"]

src/ntops/softmax.py

@@ -0,0 +1,64 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+import torch
+from ninetoothed import Tensor
+
+
+def arrangement(input, output, dim):
+    assert input.ndim == output.ndim
+
+    def create_axis_tile_shape(dim, dim_block):
+        return tuple(1 for _ in range(dim)) + (dim_block,) + tuple(1 for _ in range(input.ndim - dim - 1))
+    
+    inner_block_shape = create_axis_tile_shape(dim, input.shape[dim])
+    outer_block_shape = create_axis_tile_shape(dim, -1)
+    
+    def arrange(input):
+        input_arranged = input.tile(inner_block_shape).tile(outer_block_shape)
+
+        input_arranged.dtype = input_arranged.dtype.squeeze(
+            tuple(d for d in range(input.ndim) if d != dim)
+        )
+        input_arranged.dtype.dtype = input_arranged.dtype.dtype.squeeze(
+            tuple(d for d in range(input.ndim) if d != dim)
+        )
+        return input_arranged
+    
+    input_arranged = arrange(input)
+    output_arranged = arrange(output)
+
+    return input_arranged, output_arranged
+
+
+def application(input, output):
+    for i in range(input.shape[0]):
+        input_i = input[i]
+        row_minus_max = input_i - ntl.max(input_i)
+        numerator = ntl.exp(ntl.cast(row_minus_max, ntl.float32))
+        denominator = ntl.sum(numerator)
+        output[i] = numerator / denominator  # noqa: F841
+
+
+def softmax(input, dim, output=None):
+    if output is None:
+        output = torch.empty_like(input)
+
+    kernel = _make(input.ndim, dim)
+
+    kernel(input, output)
+
+    return output
+
+
+@functools.cache
+def _make(ndim, dim):
+    return ninetoothed.make(
+        functools.partial(arrangement, dim=dim),
+        application,
+        (
+            Tensor(ndim, other=float("-inf"), shape_options={"constexpr": True}),
+            Tensor(ndim),
+        ),
+    )

tests/test_softmax.py

@@ -0,0 +1,22 @@
+import random
+
+import pytest
+import torch
+
+import ntops
+from tests.skippers import skip_if_cuda_not_available
+from tests.utils import generate_arguments
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize(*generate_arguments())
+def test_cuda(shape, dtype, atol, rtol):
+    device = "cuda"
+
+    input = torch.randn(shape, dtype=dtype, device=device)
+    dim = random.randint(0, input.ndim - 1)
+
+    ninetoothed_output = ntops.softmax(input, dim)
+    reference_output = torch.nn.functional.softmax(input, dim=dim)
+
+    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)