softmaxVector function

Tensor<Vector> softmaxVector(Tensor<Vector> v)

Implementation

Tensor<Vector> softmaxVector(Tensor<Vector> v) {
  int N = v.data.length;
  double maxVal = -double.infinity;

  for (int i = 0; i < N; i = i + 1) {
    if (v.data[i] > maxVal) {
      maxVal = v.data[i];
    }
  }

  double sumExps = 0.0;
  Vector exps = [];
  for (int i = 0; i < N; i = i + 1) {
    double expVal = exp(v.data[i] - maxVal);
    exps.add(expVal);
    sumExps = sumExps + expVal;
  }

  Vector outValue = [];
  for (int i = 0; i < N; i = i + 1) {
    outValue.add(exps[i] / sumExps);
  }

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

  out.creator = Node(
    [v],
        () {
      double dotProduct = 0.0;
      for (int i = 0; i < N; i = i + 1) {
        dotProduct = dotProduct + (out.grad[i] * out.data[i]);
      }

      for (int i = 0; i < N; i = i + 1) {
        v.grad[i] = v.grad[i] + out.data[i] * (out.grad[i] - dotProduct);
      }
    },
    opName: 'softmax_vector',
    cost: N * 2,
  );

  return out;
}