Fix jagged NT softmax semantics (pytorch#119459)

Before: `softmax` definition uses `jagged_unary_pointwise()` (wrong) After: `softmax` impl adjusts the `dim` arg to account for the difference in dimensionality between the outer NT and the NT's `_values` Pull Request resolved: pytorch#119459 Approved by: https://github.com/soulitzer

Author: jbschlosser

test/test_nestedtensor.py

@@ -3182,6 +3182,22 @@ def test_split_with_sizes(self, device):
  ):
  torch.split(nt, [1, 2], 1)
 
+ def test_softmax(self, device):
+ nt = random_nt_from_dims(
+ [3, None, 5], device=device, dtype=torch.float32, layout=torch.jagged)
+
+ # operate on dim=2
+ output = nt.softmax(dim=2)
+ for in_component, out_component in zip(nt.unbind(), output.unbind()):
+ # dim=2 -> dim=1 after unbind
+ self.assertEqual(in_component.softmax(dim=1), out_component)
+
+ # operate on dim=-1
+ output2 = nt.softmax(dim=-1)
+ self.assertEqual(output, output2)
+ for in_component, out_component in zip(nt.unbind(), output2.unbind()):
+ self.assertEqual(in_component.softmax(dim=-1), out_component)
+
  def test_views_inherit_ragged_dim(self, device):
  # view
  nt = random_nt_from_dims(

torch/nested/_internal/ops.py

@@ -447,9 +447,19 @@ def to_copy_default(func, *args, **kwargs):
 )(jagged_unary_pointwise)
 
 
-register_jagged_func(
+@register_jagged_func(
  torch.ops.aten._softmax.default, "self: jt, dim: any, half_to_float: any"
-)(jagged_unary_pointwise)
+)
+def _softmax_default(func, *args, **kwargs):
+ _, new_kwargs = normalize_function(
+ func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+ )
+
+ inp = new_kwargs.pop("input")
+ dim = new_kwargs["dim"]
+ new_kwargs["dim"] = _wrap_jagged_dim(len(inp._size), dim, "softmax")
+
+ return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
 
 
 @register_jagged_func(