softmaxMatrix function

Tensor<Matrix> softmaxMatrix(Tensor<Matrix> m)

Mathematical operation for applying Softmax to each row of a matrix.

Implementation

Tensor<Matrix> softmaxMatrix(Tensor<Matrix> m) {
  Matrix outValue = [];
  int numRows = m.value.length;
  int numCols = m.value[0].length;

  for (int i = 0; i < numRows; i++) {
    Vector row = m.value[i];
    double maxVal = -double.infinity;
    for (double val in row) {
      if (val > maxVal) {
        maxVal = val;
      }
    }

    Vector exps = [];
    double sumExps = 0.0;
    for (double val in row) {
      double expVal = exp(val - maxVal);
      exps.add(expVal);
      sumExps += expVal;
    }

    Vector softmaxRow = [];
    for (double expVal in exps) {
      softmaxRow.add(expVal / sumExps);
    }
    outValue.add(softmaxRow);
  }

  Tensor<Matrix> out = Tensor<Matrix>(outValue);

  out.creator = Node([m], () {
    for (int r = 0; r < numRows; r++) {
      Vector softmaxRow = out.value[r];
      Vector gradRow = out.grad[r];

      double dotProduct = 0;
      for (int i = 0; i < numCols; i++) {
        dotProduct += gradRow[i] * softmaxRow[i];
      }

      for (int j = 0; j < numCols; j++) {
        m.grad[r][j] += softmaxRow[j] * (gradRow[j] - dotProduct);
      }
    }
  }, opName: 'softmax_matrix', cost: numRows * numCols * numCols);
  return out;
}