Implement more reduction

src/batch_normalization.js

@@ -20,7 +20,7 @@ export function batchNormalization(input, mean, variance, {axis=1, scale, bias,
   // The output tensor has the same shape as the input tensor.
   let output = new Tensor(input.shape);
   const shape = new Array(input.rank).fill(1);
-  shape[axis] = -1;
+  shape[axis] = null;
   output = sub(input, reshape(mean, shape));
   output = div(output,
       pow(add(reshape(variance, shape), new Scalar(epsilon)), new Scalar(0.5)));

src/gru.js

@@ -73,15 +73,15 @@ export function gru(input, weight, recurrentWeight, steps, hiddenSize,
               cellInput, cellWeight[slot], cellRecurrentWeight[slot],
               cellHidden[slot], hiddenSize, {bias: cellBias[slot],
                 recurrentBias: cellRecurrentBias[slot], resetAfter, layout, activations}),
-          [1, -1, hiddenSize]);
+          [1, null, hiddenSize]);
 
       cellOutput = (cellOutput ? concat([cellOutput, result], 0) : result);
     }
 
     hiddenState = cellOutput;
 
     if (returnSequence) {
-      cellOutput = reshape(cellOutput, [1, numDirections, -1, hiddenSize]);
+      cellOutput = reshape(cellOutput, [1, numDirections, null, hiddenSize]);
       sequence =
           (sequence ? concat([sequence, cellOutput], 0) : cellOutput);
     }

src/reduce.js

@@ -1,7 +1,9 @@
 'use strict';
 
+import {pow} from './binary.js';
 import {squeeze} from './squeeze.js';
-import {sizeOfShape, Tensor} from './lib/tensor.js';
+import {abs, exp, log} from './unary.js';
+import {sizeOfShape, Scalar, Tensor} from './lib/tensor.js';
 import {validateReduceParams} from './lib/validate-input.js';
 
 /**
@@ -16,9 +18,6 @@ function reduce(input, reduceFunc, {keepDimensions = false, axes} = {}) {
 
   const outputShape = input.shape.slice();
   for (let i = 0; i < inpAxes.length; ++i) {
-    if (inpAxes[i] === -1) {
-      inpAxes[i] = input.rank - 1;
-    }
     outputShape[inpAxes[i]] = 1;
   }
 
@@ -123,3 +122,60 @@ export function reduceSum(input, options = {}) {
   return reduce(input,
       (previousValue, currentValue) => previousValue + currentValue, options);
 }
+
+/**
+ * Compute the sum of the square of all the input values along the axes.
+ * @param {Tensor} input
+ * @param {MLReduceOptions} options
+ * @return {Tensor}
+ */
+export function reduceSumSquare(input, options = {}) {
+  return reduceSum(pow(input, new Scalar(2)), options);
+}
+
+/**
+ * Compute the L1 norm of all the input values along the axes.
+ * @param {Tensor} input
+ * @param {MLReduceOptions} options
+ * @return {Tensor}
+ */
+export function reduceL1(input, options = {}) {
+  return reduceSum(abs(input), options);
+}
+
+/**
+ * Compute the L2 norm of all the input values along the axes.
+ * @param {Tensor} input
+ * @param {MLReduceOptions} options
+ * @return {Tensor}
+ */
+export function reduceL2(input, options = {}) {
+  const intermediateResult = reduceSumSquare(input, options);
+  if (intermediateResult.rank === 0) {
+    return new Tensor(
+        [],
+        [Math.pow(intermediateResult.getValueByIndex(0), 0.5)]);
+  } else {
+    return pow(intermediateResult, new Scalar(0.5));
+  }
+}
+
+/**
+ * Compute the log value of the sum of all the input values along the axes.
+ * @param {Tensor} input
+ * @param {MLReduceOptions} options
+ * @return {Tensor}
+ */
+export function reduceLogSum(input, options = {}) {
+  return log(reduceSum(input, options));
+}
+
+/**
+ * Compute the log value of the sum of the exponent of all the input values along the axes.
+ * @param {Tensor} input
+ * @param {MLReduceOptions} options
+ * @return {Tensor}
+ */
+export function reduceLogSumExp(input, options = {}) {
+  return log(reduceSum(exp(input), options));
+}

src/reshape.js

@@ -12,7 +12,7 @@ export function reshape(input, newShape) {
   let minusOneAxis;
   let elements = 1;
   for (let i = 0; i < newShape.length; ++i) {
-    if (newShape[i] === -1) {
+    if (newShape[i] === null) {
       minusOneAxis = i;
     } else if (newShape[i] > 0) {
       elements *= newShape[i];

src/slice.js

@@ -20,7 +20,7 @@ export function slice(input, starts, sizes, {axes} = {}) {
   const axesLen = axes.length;
   const outputShape = input.shape.slice();
   for (let i = 0; i < axesLen; ++i) {
-    const axis = axes[i] >= 0 ? axes[i] : axes[i] + rank;
+    const axis = axes[i];
     const size = input.shape[axis];
     const start = starts[i];
     startsForAllAxes[axis] = start >= 0 ? start : start + size;

src/split.js

@@ -14,16 +14,14 @@ export function split(input, splits, {axis = 0} = {}) {
   validateSplitParams(...arguments);
   const outputs = [];
   let sliceSizes = [];
-  const rank = input.rank;
-  const inpAxis = axis >=0 ? axis : rank + axis;
   if (typeof splits === 'number') {
-    sliceSizes = new Array(splits).fill(input.shape[inpAxis] / splits);
+    sliceSizes = new Array(splits).fill(input.shape[axis] / splits);
   } else if (splits instanceof Array) {
     sliceSizes = splits.slice();
   }
   let start = 0;
   for (const size of sliceSizes) {
-    outputs.push(slice(input, [start], [size], {axes: [inpAxis]}));
+    outputs.push(slice(input, [start], [size], {axes: [axis]}));
     start += size;
   }
   return outputs;