Why tensorflow use 'dim' parameter for softmax function?

Question

Why tensorflow use 'dim' parameter for softmax function? What kind of tensors we can use as input ?

Answer 1

tf.nn.softmax accepts in input a generic nonempty tensor.

You can decide to apply softmax on every dimension you want to.

Usually, softmax is applied to the last dimension (that's the default behavior) of the input tensor. This because usually softmax is applied to neural network output that's usually a tensor with a shape of [batch_size, num_classes].

However, you could decide to apply softmax to a tensor with a shape of [batch_size, num_classes, 2, 1] and compute the softmax only over the second dimension of the tensor: tf.nn.softmax(tensor, axis=1)