tf.sparse.softmax  |  TensorFlow v2.16.1 (original) (raw)

tf.sparse.softmax

Stay organized with collections Save and categorize content based on your preferences.

Applies softmax to a batched N-D SparseTensor.

View aliases

Compat aliases for migration

SeeMigration guide for more details.

tf.compat.v1.sparse.softmax, tf.compat.v1.sparse_softmax

tf.sparse.softmax(
    sp_input, name=None
)

The inputs represent an N-D SparseTensor with logical shape [..., B, C](where N >= 2), and with indices sorted in the canonical lexicographic order.

This op is equivalent to applying the normal tf.nn.softmax() to each innermost logical submatrix with shape [B, C], but with the catch that the implicitly zero elements do not participate. Specifically, the algorithm is equivalent to:

(1) Applies tf.nn.softmax() to a densified view of each innermost submatrix with shape [B, C], along the size-C dimension; (2) Masks out the original implicitly-zero locations; (3) Renormalizes the remaining elements.

Hence, the SparseTensor result has exactly the same non-zero indices and shape.

Example using a 3-D SparseTensor:

st = tf.sparse.from_dense( [[[0., np.e], [1., 0.]], `` [[np.e, 0.], [np.e, np.e]]]) res = tf.sparse.softmax(st) res.indices <tf.Tensor: shape=(5, 3), dtype=int64, numpy= array([[0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 1, 0], [1, 1, 1]])> res.values <tf.Tensor: ... numpy=array([1. , 1. , 1. , 0.5, 0.5], dtype=float32)> res.dense_shape <tf.Tensor: shape=(3,), dtype=int64, numpy=array([2, 2, 2])> tf.sparse.to_dense(res) <tf.Tensor: shape=(2, 2, 2), dtype=float32, numpy= array([[[0. , 1. ], [1. , 0. ]], [[1. , 0. ], [0.5, 0.5]]], dtype=float32)>