softmax method

Tensor softmax(
  1. int dim
)

Softmax over a given dimension.

Implementation

Tensor softmax(int dim) {
  if (dim < 0) dim = ndim + dim;
  final maxVals = maxDim(dim, keepDim: true);
  final shifted = this - maxVals.expand(shape);
  final exps = shifted.map((x) => math.exp(x));
  final sumExps = exps.sum(dim, keepDim: true);
  return exps / sumExps.expand(shape);
}