The input array.
The axis along which to compute softmax.
DType of the output in case this can't be inferred. Defaults to the same as input's dtype if not defined (dtype=None).
Temperature parameter in softmax
Whether to use the length input as a mask over the data input.