maxDim method

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

Max over a given dimension.

Implementation

Tensor maxDim(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 m = data[outer * dimSize * innerSize + inner];
      for (int d = 1; d < dimSize; d++) {
        final v = data[outer * dimSize * innerSize + d * innerSize + inner];
        if (v > m) m = v;
      }
      result.data[outer * innerSize + inner] = m;
    }
  }
  if (!keepDim) {
    return result.squeeze(dim);
  }
  return result;
}