add cumulativeSum

src/cumulativeSum.js

@@ -0,0 +1,46 @@
+'use strict';
+
+import {Tensor, sizeOfShape} from './lib/tensor.js';
+import {validateCumulativeSumParams} from './lib/validate-input.js';
+
+/**
+ * Computes the cumulative sum of the input tensor along the specified axis.
+ * @param {Tensor} input
+ * @param {number} axis
+ * @param {MLCumulativeSumOptions} options
+ * @return {Tensor}
+ */
+export function cumulativeSum(input, axis, {exclusive = 0, reverse = 0} = {}) {
+  validateCumulativeSumParams(...arguments);
+  const inputShape = input.shape;
+  const outputShape = [...inputShape];
+  const output = new Tensor(outputShape);
+  const numElementsAlongAxis = inputShape[axis];
+
+  const totalElements = sizeOfShape(outputShape);
+
+  for (let outputIndex = 0; outputIndex < totalElements; outputIndex++) {
+    const loc = output.locationFromIndex(outputIndex);
+    let cumulativeSumValue = 0;
+
+    const start = reverse ? numElementsAlongAxis - 1 : 0;
+    const step = reverse ? -1 : 1;
+    const end = reverse ? -1 : numElementsAlongAxis;
+
+    for (let i = start; reverse ? i > end : i < end; i += step) {
+      const inputLoc = [...loc];
+      inputLoc[axis] = exclusive ? (reverse ? i + 1 : i - 1) : i;
+
+      if (!exclusive || (exclusive && inputLoc[axis] >= 0 &&
+        inputLoc[axis] < numElementsAlongAxis)) {
+        cumulativeSumValue += input.getValueByLocation(inputLoc);
+      }
+
+      const outputLoc = [...loc];
+      outputLoc[axis] = i;
+      output.setValueByLocation(outputLoc, cumulativeSumValue);
+    }
+  }
+
+  return output;
+}

src/lib/validate-input.js

@@ -611,6 +611,18 @@ export function validateGatherParams(input, indices, {axis = 0} = {}) {
   }
 }
 
+export function validateCumulativeSumParams(input, axis) {
+  if (axis !== undefined) {
+    const rank = input.rank;
+    if (!Number.isInteger(axis) || axis < -rank || axis >= rank) {
+      throw new Error(`The axis ${axis} should be in the range [-rank(input), rank(input)-1].`);
+    }
+    if (axis >= rank) {
+      throw new Error(`The axis ${axis} should be in the interval [0, ${rank}).`);
+    }
+  }
+}
+
 export function validateTriangularParams(input, {diagonal = 0} = {}) {
   const inputRank = input.rank;
   if (inputRank < 2) {

test/cumulativeSum_test.js

@@ -0,0 +1,143 @@
+'use strict';
+
+import {cumulativeSum} from '../src/cumulativeSum.js';
+import {Tensor} from '../src/lib/tensor.js';
+import * as utils from './utils.js';
+
+describe('test cumulativeSum', function() {
+  function testCumulativeSum(input, axis, options={}, expected) {
+    const tensor = new Tensor(input.shape, input.data);
+    const outputTensor = cumulativeSum(tensor, axis, options);
+    console.log('outputTensor', outputTensor);
+    utils.checkShape(outputTensor, expected.shape);
+    utils.checkValue(outputTensor, expected.data);
+  }
+
+  it('test cumulativeSum 1d', function() {
+    const input = {
+      shape: [5],
+      data: [
+        1, 2, 3, 4, 5,
+      ],
+    };
+    const axis=0;
+    const options = {exclusive: 0, reverse: 0};
+    const expected = {
+      shape: [5],
+      data: [
+        1, 3, 6, 10, 15,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+
+  it('test cumulativeSum 1d exclusive', function() {
+    const input = {
+      shape: [5],
+      data: [
+        1, 2, 3, 4, 5,
+      ],
+    };
+    const axis=0;
+    const options = {exclusive: 1, reverse: 0};
+    const expected = {
+      shape: [5],
+      data: [
+        0, 1, 3, 6, 10,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+
+  it('test cumulativeSum 1d reverse', function() {
+    const input = {
+      shape: [5],
+      data: [
+        1, 2, 3, 4, 5,
+      ],
+    };
+    const axis=0;
+    const options = {exclusive: 0, reverse: 1};
+    const expected = {
+      shape: [5],
+      data: [
+        15, 14, 12, 9, 5,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+
+  it('test cumulativeSum 1d reverse exclusive', function() {
+    const input = {
+      shape: [5],
+      data: [
+        1, 2, 3, 4, 5,
+      ],
+    };
+    const axis=0;
+    const options = {exclusive: 1, reverse: 1};
+    const expected = {
+      shape: [5],
+      data: [
+        14, 12, 9, 5, 0,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+
+  it('test cumulativeSum 2d', function() {
+    const input = {
+      shape: [2, 3],
+      data: [
+        1, 2, 3, 4, 5, 6,
+      ],
+    };
+    const axis=0;
+    const options = {exclusive: 0, reverse: 0};
+    const expected = {
+      shape: [2, 3],
+      data: [
+        1, 2, 3, 5, 7, 9,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+
+  it('test cumulativeSum 2d axis=1', function() {
+    const input = {
+      shape: [2, 3],
+      data: [
+        1, 2, 3, 4, 5, 6,
+      ],
+    };
+    const axis=1;
+    const options = {exclusive: 0, reverse: 0};
+    const expected = {
+      shape: [2, 3],
+      data: [
+        1, 3, 6, 4, 9, 15,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+
+  it('test cumulativeSum 2d negtive axis', function() {
+    const input = {
+      shape: [2, 3],
+      data: [
+        1, 2, 3, 4, 5, 6,
+      ],
+    };
+    const axis=1;
+    const options = {exclusive: 0, reverse: 0};
+    const expected = {
+      shape: [2, 3],
+      data: [
+        1, 3, 6, 4, 9, 15,
+      ],
+    };
+    testCumulativeSum(input, axis, options, expected);
+  });
+});
+
+