sum method

Tensor sum(int dim, { bool keepDim = false, })

Sum over a given dimension, keeping the dimension (size 1).

Implementation

Tensor sum(int dim, {bool keepDim = false}) {
  if (dim < 0) dim = ndim + dim;
  final outShape = List<int>.from(shape);
  outShape[dim] = 1;
  final result = Tensor.zeros(outShape);

  final outerSize = _productOfShape(shape.sublist(0, dim));
  final dimSize = shape[dim];
  final innerSize =
      dim + 1 < ndim ? _productOfShape(shape.sublist(dim + 1)) : 1;

  for (int outer = 0; outer < outerSize; outer++) {
    for (int inner = 0; inner < innerSize; inner++) {
      double s = 0.0;
      for (int d = 0; d < dimSize; d++) {
        s += data[outer * dimSize * innerSize + d * innerSize + inner];
      }
      result.data[outer * innerSize + inner] = s;
    }
  }
  if (!keepDim) {
    return result.squeeze(dim);
  }
  return result;
}