implement gather

src/gather.js

@@ -0,0 +1,60 @@
+'use strict';
+
+import {Tensor, sizeOfShape} from './lib/tensor.js';
+import {validateGatherParams} from './lib/validate-input.js';
+
+/**
+ * Gather values of the input tensor along an axis according to the indices.
+ * @param {Tensor} input
+ * @param {Tensor} indices
+ * @param {MLGatherOptions} [options]
+ * @return {Tensor}
+ */
+export function gather(input, indices, {axis = 0} = {}) {
+  validateGatherParams(...arguments);
+  const shapeInput = input.shape;
+
+  // set outputShape following Spec Algorithm
+  //   https://webmachinelearning.github.io/webnn/#api-mlgraphbuilder-gather
+  //
+  // let dimCount = 0
+  // let rankOutput = 0;
+  // let shapeOutput = [];
+  // for (dimCount = 0; dimCount < shapeInput.length; dimCount++) {
+  //   if (dimCount === axis) {
+  //     break;
+  //   } else {
+  //     shapeOutput[dimCount] = shapeInput[dimCount];
+  //   }
+  // }
+  // rankOutput = dimCount;
+  // for (dimCount = 0; dimCount < indices.shape.length; dimCount++) {
+  //   shapeOutput[rankOutput + dimCount] = indices.shape[dimCount];
+  // }
+  // rankOutput = rankOutput + dimCount;
+  // for (dimCount = 0; dimCount < shapeInput.length; dimCount++) {
+  //   if (dimCount <= axis) {
+  //     continue;
+  //   } else {
+  //     shapeOutput[rankOutput + dimCount - axis - 1] = shapeInput[dimCount];
+  //   }
+  // }
+
+  // optimized set outputShape using JavaScipt slice and concat
+  const shapeOutput = shapeInput.slice(0, axis).concat(indices.shape, shapeInput.slice(axis + 1));
+  const output = new Tensor(shapeOutput);
+
+  for (let outputIndex = 0; outputIndex < sizeOfShape(shapeOutput); ++outputIndex) {
+    // output[i, j, k, ...] = input[indices[i, j, k, ...], j, k, ...] // if axis == 0
+    // output[i, j, k, ...] = input[i, indices[i, j, k, ...], k, ...] // if axis == 1
+    // output[i, j, k, ...] = input[i, j, indices[i, j, k, ...], ...] // if axis == 2
+    const outputLoc = output.locationFromIndex(outputIndex);
+    const indicesLoc = outputLoc.slice(axis, axis + indices.rank);
+    const selectedInputLoc = outputLoc.slice(0, axis)
+        .concat(indices.getValueByLocation(indicesLoc), outputLoc.slice(axis + indices.rank));
+    const inputValue = input.getValueByLocation(selectedInputLoc);
+    output.setValueByIndex(outputIndex, inputValue);
+  }
+
+  return output;
+}

src/lib/validate-input.js

@@ -390,3 +390,23 @@ export function validateNotParams(input) {
     }
   }
 }
+
+export function validateGatherParams(input, indices, {axis = 0} = {}) {
+  if (axis !== undefined) {
+    const rank = input.rank;
+    if (!Number.isInteger(axis) || axis < 0) {
+      throw new Error(`The axis ${axis} should be an unsigned integer.`);
+    }
+    if (axis >= rank) {
+      throw new Error(`The axis ${axis} should be in the interval [0, ${rank}).`);
+    }
+  }
+  const axisSize = input.shape[axis];
+  for (let i = 0; i < sizeOfShape(indices.shape); ++i) {
+    const index = indices.getValueByIndex(i);
+    if (!Number.isInteger(index) || index < 0 || index >= axisSize) {
+      throw new Error(
+          `Invalid indices value - it should be an integer in the interval [0, ${axisSize})`);
+    }
+  }
+}