math icon indicating copy to clipboard operation
math copied to clipboard

Add additional signatures to `softmax` and `log_softmax`

Open jachymb opened this issue 9 months ago • 0 comments

Currently the only signature available is (vector) => vector

But it's somewhat annoying when I want to do this row-wise over a matrix. (To hand-code a similar effect as stochastic_row_matrix constraint has)

I suggest adding the following signatures:

Softmax for row_vectors:

(row_vector) => row_vector

And a 'vectorized' variant:

(array[] vector) => array[] vector (array[] row_vector) => array[] row_vector

The signature (matrix) => matrix is perhaps not a good idea because it's not immediately clear which dimension it should work over, but maybe there could be pairs of functions softmax_rowwise, softmax_columnwise (which would essentially return stochastic_row_matrix and stochastic_col_matrix respectively) as well as the log variants log_softmax_rowwise, log_softmax_colwise.

Possibly related: https://github.com/stan-dev/stan/issues/3331

jachymb avatar Mar 19 '25 07:03 jachymb