The sparse softmax is related to dense softmax when the values of unspecified sparse tensor entries are taken to be -inf that will have the effect of "zero entries ignored"

We had a bunch of discussions about whether or not it is appropriate to treat unspecified sparse tensor entries as things other than zero; specifically, whether or not we should be allowed to overload the old function name in this situation (some discussion about this at #1369 ). Did you folks at Quansight come to a decision about what should be done here?

@ezyang we have not discussed the overloading vs name-prefixing options yet but we have discussed the idea of introducing the so-called fill value to sparse tensors that would set the default value for unspecified sparse tensor entries.

The fill value feature would be useful also for the sparse softmax case that would allow users to specify the user-defined behavior:

• Using fill_value=0 would lead to nice invariant softmax(sparse, dim).to_dense() == softmax(sparse.to_dense(), dim) while the result of softmax(sparse, dim) would be still sparse tensor (with a non-zero fill_value).
• Using fill_value=-inf [the current PR] gives a mathematical basis for softmax(sparse, dim) to be a sparse tensor (as requested in the issue Add sparse softmax/log_softmax functionality (ignore zero entries) #23651).

In addition, the fill value feature has the potential to eliminate the overloading vs name-prefixing question, and in general, it would lower the bridge between sparse and dense tensor in a consistent way.

The fill value has been implemented in https://sparse.pydata.org/ and from #9674 (comment) I see you have already discussed this option as well.

For this PR, we have two choices:

1. sparse softmax is exposed as softmax (and log_softmax) [the current behavior] via overloading
2. sparse softmax is exposed as spsoftmax (and splog_softmax)

In the long term, with the perspective of having the fill value feature available, option 1. makes sense to me.
On the other hand, option 2. seems to be the current practice in pytorch and I would be ok changing the PR accordingly.

So, I am not certain atm which approach would be preferable to users. Thoughts?

If y'all are prepared to take on the maintenance burden of a fill value, I think I would also agree that this would be best, and then we don't have to introduce sparse versions of all of our functions. Without fill, (1) will technically require you to do BC-breaking changes in the future, as you will turn an operation that previously implicitly converted zero-filled tensors to neginf-filled tensors, into an error. It will save you work in the future if you do (2) first.

OK, it makes sense. Here's a plan:

• Expose sparse softmax as torch.sparse.softmax that is defined for sparse tensors with assumed fill value equal to negative infinity (that would be equivalent to ignoring unspecified sparse tensor entries).
• Expose sparse log softmax as torch.sparse.log_softmax with fill value as above.
• Defining torch.softmax on sparse tensors will be reserved for sparse tensors that support fill value feature. The result of torch.softmax will be different from the result of torch.sparse.softmax because of the default fill value for sparse tensors will be zero. The results will be equal iff the fill value of a sparse tensor is specified as negative infinity.

SGTM.

@apaszke @ezyang this PR is ready for review. The PR implements sparse softmax support for CPU, the CUDA support is planned as a followup PR.

@pearu can you get someone from Quansight to review first?

All the loops could be replaced with the TensorIterator loops, I think... TensorIterator loops over all dimensions not involved in reduction, and the kernel loop iterates over the reduction dimension if you perform the reduction and scaling all at once. Otherwise you could iterate over each element and then rescale it appropriately while knowing which value to rescale it with. TensorIterator also brings multithreading for free.

Here are some benchmark results for torch.sparse.softmax (v1.6.0.a) and tensorflow.sparse.softmax (v 2.2.0), both using CPU, values are float64 scalars, timing repeat is 1000:

Some notes:

• pytorch.sparse.softmax is generally faster than tensorflow.sparse.softmax. That is pretty good considering that pytorch sparse softmax supports more features that tf.
• softmax performance is higher for larger dimension argument, this is likely related to better locality properties of index pools when using row-wise storage.
• pytorch sparse vs dense timings show that there exists a density region (< 1 %) where using sparse storage gives considerably better performance than using dense storage. For relatively dense matrices, usage of sparse storage leads to up to 10x smaller performance.

Failure is legit:

May 12 18:18:16   test_softmax (__main__.TestSparse) ... /var/lib/jenkins/workspace/aten/src/ATen/native/sparse/SoftMax.cpp:348:35: runtime error: division by zero
May 12 18:18:17     #0 0x7f50670b37cd in .omp_outlined. (/opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_cpu.so+0xc2517cd)
May 12 18:18:17     #1 0x7f5053405712 in __kmp_invoke_microtask (/opt/conda/lib/libiomp5.so+0x122712)
May 12 18:18:17     #2 0x7f5053393ffe in __kmp_invoke_task_func (/opt/conda/lib/libiomp5.so+0xb0ffe)
May 12 18:18:17     #3 0x7f5053393059 in __kmp_launch_thread (/opt/conda/lib/libiomp5.so+0xb0059)
May 12 18:18:17     #4 0x7f5053405bd7 in _INTERNALdc942ded::__kmp_launch_worker(void*) (/opt/conda/lib/libiomp5.so+0x122bd7)
May 12 18:18:17     #5 0x7f508a9cd6b9 in start_thread (/lib/x86_64-linux-gnu/libpthread.so.0+0x76b9)
May 12 18:18:17     #6 0x7f508a70341c in clone /build/glibc-LK5gWL/glibc-2.23/misc/../sysdeps/unix/sysv/linux/x86_64/clone.S:109
May 12 18:18:17 

@ezyang merged this pull request in 48c0331.

Interesting to read the conversations in this PR again in light of maskedtensor cc @george-qi