optimized axis validation and outputLocation setting

webmachinelearning · Oct 8, 2024 · 8b20da9 · 8b20da9
1 parent d0807b9
commit 8b20da9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 9 deletions.
diff --git a/src/lib/validate-input.js b/src/lib/validate-input.js
@@ -633,11 +633,9 @@ export function validateScatterElementsParams(input, indices, updates, {axis = 0
   if (!updates.shape.every((value, index) => value === indices.shape[index])) {
     throw new Error(`Invalid updates value - updates shape should be same as indices shape.`);
   }
-  if (axis !== undefined) {
-    if (!Number.isInteger(axis) || axis < 0 || axis >= inputRank) {
-      throw new Error(
-          `The axis ${axis} should be an unsigned integer in the interval [0, ${inputRank}).`);
-    }
+  if (!Number.isInteger(axis) || axis < 0 || axis >= inputRank) {
+    throw new Error(
+        `The axis ${axis} should be an unsigned integer in the interval [0, ${inputRank}).`);
   }
   const axisSize = input.shape[axis];
   for (let i = 0; i < sizeOfShape(indices.shape); ++i) {

diff --git a/src/scatter_elements.js b/src/scatter_elements.js
@@ -22,10 +22,10 @@ export function scatterElements(input, indices, updates, {axis = 0} = {}) {
     // output[i, indices[i, j, k, ...], k, ...] = updates[i, j, k, ...] // if axis == 1
     // output[i, j, indices[i, j, k, ...], ...] = updates[i, j, k, ...] // if axis == 2
     const indicesLocation = indices.locationFromIndex(indicesIndex);
-    let indiceValue = indices.getValueByIndex(indicesIndex);
-    indiceValue = indiceValue < 0 ? indiceValue + input.shape[axis] : indiceValue;
-    const outputLocation = indicesLocation.slice();
-    outputLocation[axis] = indiceValue;
+    let indicesValue = indices.getValueByIndex(indicesIndex);
+    indicesValue = indicesValue < 0 ? indicesValue + input.shape[axis] : indicesValue;
+    const outputLocation =
+        [...indicesLocation.slice(0, axis), indicesValue, ...indicesLocation .slice(axis + 1)];
     output.setValueByLocation(outputLocation, updates.getValueByIndex(indicesIndex));
   }