[tfjs-core] Support newAxisMask in stridedSlice and make squeez… (#1829)

FEATURE
BUG

* Add support for the `newAxisMask` param in `stridedSlice()`. Because stridedSlice() is difficult to implement correctly, and all backends have to repeat fragile logic, shift common preprocess and postprocess logic to the logical stridedSlice op and simplify the actual kernel.
* Make tf.squeeze(x, axis) consistent with TF when axis is empty array.

These fixes were found by running a converted TF model in the browser.

Author: dsmilkov

tfjs-core/src/backends/backend.ts

@@ -141,9 +141,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
     throw new Error('Not yet implemented');
   }
   stridedSlice<T extends Tensor>(
-      x: T, begin: number[], end: number[], strides: number[],
-      beginMask: number, endMask: number, ellipsisMask: number,
-      newAxisMask: number, shrinkAxisMask: number): T {
+      x: T, begin: number[], end: number[], strides: number[]): T {
     throw new Error('Not yet implemented');
   }
   unstack(x: Tensor, axis: number): Tensor[] {

tfjs-core/src/backends/cpu/backend_cpu.ts

@@ -33,7 +33,7 @@ import * as ops from '../../ops/ops';
 import {buffer, scalar, tensor, tensor3d, tensor4d} from '../../ops/ops';
 import * as scatter_nd_util from '../../ops/scatter_nd_util';
 import * as selu_util from '../../ops/selu_util';
-import {computeFlatOffset, getStridedSlicedInfo, isSliceContinous} from '../../ops/slice_util';
+import {computeFlatOffset, computeOutShape, isSliceContinous} from '../../ops/slice_util';
 import {DataId, Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer} from '../../tensor';
 import {BackendValues, DataType, DataValues, NumericDataType, PixelData, Rank, ShapeMap, TypedArray, upcastType} from '../../types';
 import * as util from '../../util';
@@ -313,34 +313,28 @@ export class MathBackendCPU implements KernelBackend {
   }
 
   stridedSlice<T extends Tensor>(
-      x: T, begin: number[], end: number[], strides: number[],
-      beginMask: number, endMask: number, ellipsisMask: number,
-      newAxisMask: number, shrinkAxisMask: number): T {
+      x: T, begin: number[], end: number[], strides: number[]): T {
     this.assertNotComplex(x, 'stridedSlice');
 
-    const [beginIndex, size, shrinkAxis] = getStridedSlicedInfo(
-        x.shape, begin, end, strides, beginMask, endMask, ellipsisMask,
-        newAxisMask, shrinkAxisMask);
+    const outShape = computeOutShape(begin, end, strides);
 
-    const shape = size.filter((v, index) => shrinkAxis.indexOf(index) === -1);
-
-    if (shape.some(axis => axis === 0)) {
-      return ops.tensor([], shape) as T;
+    if (outShape.some(axis => axis === 0)) {
+      return ops.tensor([], outShape) as T;
     }
 
-    const buffer = ops.buffer(size, x.dtype);
+    const buffer = ops.buffer(outShape, x.dtype);
     const xBuf = this.bufferSync(x);
     for (let i = 0; i < buffer.size; i++) {
       const loc = buffer.indexToLoc(i);
 
       const newLoc: number[] = new Array(loc.length);
       for (let j = 0; j < newLoc.length; j++) {
-        newLoc[j] = loc[j] * strides[j] + beginIndex[j];
+        newLoc[j] = loc[j] * strides[j] + begin[j];
       }
       buffer.set(xBuf.get(...newLoc), ...loc);
     }
 
-    return buffer.toTensor().reshape(shape) as T;
+    return buffer.toTensor() as T;
   }
 
   diag(x: Tensor): Tensor {

tfjs-core/src/backends/webgl/backend_webgl.ts

@@ -34,7 +34,7 @@ import * as gather_nd_util from '../../ops/gather_nd_util';
 import * as reduce_util from '../../ops/reduce_util';
 import * as scatter_nd_util from '../../ops/scatter_nd_util';
 import * as segment_util from '../../ops/segment_util';
-import {computeFlatOffset, getStridedSlicedInfo, isSliceContinous} from '../../ops/slice_util';
+import * as slice_util from '../../ops/slice_util';
 import {softmax} from '../../ops/softmax';
 import {range, scalar, tensor} from '../../ops/tensor_ops';
 import {DataId, Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D} from '../../tensor';
@@ -729,7 +729,7 @@ export class MathBackendWebGL implements KernelBackend {
       return tensor([], size, x.dtype) as T;
     }
     const {isPacked} = this.texData.get(x.dataId);
-    const isContinous = isSliceContinous(x.shape, begin, size);
+    const isContinous = slice_util.isSliceContinous(x.shape, begin, size);
     if (isPacked || !isContinous) {
       const program = ENV.getBool('WEBGL_PACK_ARRAY_OPERATIONS') ?
           new SlicePackedProgram(size) :
@@ -749,7 +749,7 @@ export class MathBackendWebGL implements KernelBackend {
     Object.assign(newTexData, xTexData);
     newTexData.shape = size;
     newTexData.dtype = x.dtype;
-    let flatOffset = computeFlatOffset(begin, x.strides);
+    let flatOffset = slice_util.computeFlatOffset(begin, x.strides);
     if (xTexData.slice) {
       // We are slicing an already sliced tensor, so we have to accumulate
       // the offset.
@@ -769,26 +769,18 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   stridedSlice<T extends Tensor>(
-      x: T, begin: number[], end: number[], strides: number[],
-      beginMask: number, endMask: number, ellipsisMask: number,
-      newAxisMask: number, shrinkAxisMask: number): T {
+      x: T, begin: number[], end: number[], strides: number[]): T {
     if (this.shouldExecuteOnCPU([x])) {
-      return this.cpuBackend.stridedSlice(
-          x, begin, end, strides, beginMask, endMask, ellipsisMask, newAxisMask,
-          shrinkAxisMask);
+      return this.cpuBackend.stridedSlice(x, begin, end, strides);
     }
 
-    const [beginIndex, size, shrinkAxis] = getStridedSlicedInfo(
-        x.shape, begin, end, strides, beginMask, endMask, ellipsisMask,
-        newAxisMask, shrinkAxisMask);
+    const outShape = slice_util.computeOutShape(begin, end, strides);
 
-    const shape = size.filter((v, index) => shrinkAxis.indexOf(index) === -1);
-    if (shape.some(axis => axis === 0)) {
-      return tensor([], shape) as T;
+    if (outShape.some(axis => axis === 0)) {
+      return tensor([], outShape) as T;
     }
 
-    const program =
-        new StridedSliceProgram(beginIndex, strides, size, shrinkAxis);
+    const program = new StridedSliceProgram(begin, strides, outShape);
     return this.compileAndRun(program, [x]);
   }
 

tfjs-core/src/backends/webgl/strided_slice_gpu.ts

@@ -24,13 +24,11 @@ export class StridedSliceProgram implements GPGPUProgram {
   userCode: string;
 
   constructor(
-      begin: number[], strides: number[], size: number[],
-      shrinkAxis: number[]) {
-    const shape = size.filter((v, index) => shrinkAxis.indexOf(index) === -1);
-    this.outputShape = shape;
+      begin: number[], strides: number[], size: number[]) {
+    this.outputShape = size;
     const rank = size.length;
     const inputDtype = getCoordsDataType(size.length);
-    const dtype = getCoordsDataType(shape.length);
+    const dtype = getCoordsDataType(size.length);
 
     let newCoords = '';
     if (rank === 1) {
@@ -39,14 +37,10 @@ export class StridedSliceProgram implements GPGPUProgram {
       let outputAxis = 0;
       newCoords =
           size.map((_, i) => {
-                if (shrinkAxis.indexOf(i) === -1) {
-                  outputAxis++;
-                  return shape.length === 1 ?
-                      `coords * strides[${i}] + begin[${i}]` :
-                      `coords[${outputAxis - 1}] * strides[${i}] + begin[${i}]`;
-                } else {
-                  return `begin[${i}]`;
-                }
+                outputAxis++;
+                return size.length === 1 ?
+                    `coords * strides[${i}] + begin[${i}]` :
+                    `coords[${outputAxis - 1}] * strides[${i}] + begin[${i}]`;
               })
               .join(',');
     }

tfjs-core/src/ops/slice_util.ts

@@ -38,49 +38,28 @@ export function assertParamsValid(
   }
 }
 
-/**
- * Calculate the start index and output tensor shape for strided slice op.
- * @returns array of [startIndex, size, shrinkAxis]
- */
-export function getStridedSlicedInfo(
-    shape: number[], begin: number[], end: number[], strides: number[],
-    beginMask = 0, endMask = 0, ellipsisMask = 0, newAxisMask = 0,
-    shrinkAxisMask = 0): [number[], number[], number[]] {
-  if (ellipsisMask !== 0) {
-    throw new Error('ellipsis mask is not yet supported');
-  }
-  if (newAxisMask !== 0) {
-    throw new Error('new axis mask is not yet supported');
-  }
-  // Note that the axis orders are reversed for runtime ops, so the indices,
-  // strides and masks must be as well too.
-  const startIndex: number[] = [];
-  const endIndex: number[] = [];
-  const shrinkAxis: number[] = [];
-  for (let i = 0; i < shape.length; i++) {
-    startIndex[i] = startForAxis(beginMask, begin, strides, shape, i);
-    endIndex[i] = stopForAxis(endMask, end, strides, shape, i);
-    // When shrinking an axis, use startIndex + 1 for endIndex.
-    // Check the axis bit from right of shrinkAxisMask
-    if (shrinkAxisMask & 1 << i) {
-      endIndex[i] = startIndex[i] + 1;
-      shrinkAxis.push(i);
+/** Converts a binary mask to an array of axes. Used in stridedSlice(). */
+export function maskToAxes(mask: number): number[] {
+  const axes = [];
+  let axis = 0;
+  while (mask > 0) {
+    if (mask & 1) {
+      axes.push(axis);
     }
+    mask /= 2;
+    axis++;
   }
+  return axes;
+}
 
-  let size = new Array(shape.length).fill(0);
-  size = size.map((d, i) => {
-    let count = 0;
-    const stride = strides[i] || 1;
-    for (let start = startIndex[i];
-         !(stride > 0 ? start >= endIndex[i] : start <= endIndex[i]);
-         start += stride) {
-      count += 1;
-    }
-    return count;
-  });
-
-  return [startIndex, size, shrinkAxis];
+/** Computes the output shape given the strided slice params. */
+export function computeOutShape(
+    begin: number[], end: number[], strides: number[]): number[] {
+  const size = [];
+  for (let axis = 0; axis < begin.length; axis++) {
+    size[axis] = Math.ceil((end[axis] - begin[axis]) / strides[axis]);
+  }
+  return size;
 }
 
 export function startForAxis(

tfjs-core/src/ops/strided_slice.ts

@@ -19,9 +19,10 @@ import {ENGINE} from '../engine';
 import {Tensor} from '../tensor';
 import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
+
 import {op} from './operation';
 import {slice} from './slice';
-import {getStridedSlicedInfo} from './slice_util';
+import {computeOutShape, maskToAxes, startForAxis, stopForAxis} from './slice_util';
 
 /**
  * Extracts a strided slice of a tensor.
@@ -56,30 +57,53 @@ import {getStridedSlicedInfo} from './slice_util';
  */
 /** @doc {heading: 'Operations', subheading: 'Slicing and Joining'} */
 function stridedSlice_(
-    x: Tensor|TensorLike, begin: number[], end: number[], strides: number[],
+    x: Tensor|TensorLike, begin: number[], end: number[], strides?: number[],
     beginMask = 0, endMask = 0, ellipsisMask = 0, newAxisMask = 0,
     shrinkAxisMask = 0): Tensor {
+  if (strides == null) {
+    strides = new Array(begin.length);
+  }
   if (ellipsisMask !== 0) {
     throw new Error('ellipsis mask is not yet supported');
   }
-  if (newAxisMask !== 0) {
-    throw new Error('new axis mask is not yet supported');
+  let $x = convertToTensor(x, 'x', 'stridedSlice');
+
+  // Expand the dims of x based on the newAxisMask.
+  const expandAxes = maskToAxes(newAxisMask);
+  const newShape = $x.shape.slice();
+  expandAxes.forEach(axis => {
+    begin[axis] = 0;
+    end[axis] = 1;
+    newShape.splice(axis, 0, 1);
+  });
+  $x = $x.reshape(newShape);
+
+  // Normalize the start, end and strides.
+  for (let axis = 0; axis < $x.rank; axis++) {
+    begin[axis] = startForAxis(beginMask, begin, strides, $x.shape, axis);
+    end[axis] = stopForAxis(endMask, end, strides, $x.shape, axis);
+    strides[axis] = strides[axis] || 1;
   }
-  const $x = convertToTensor(x, 'x', 'stridedSlice');
+
+  const shrinkAxes = maskToAxes(shrinkAxisMask);
+  // Adjust the ends based on the shrink mask.
+  shrinkAxes.forEach(axis => {
+    end[axis] = begin[axis] + 1;
+    strides[axis] = 1;
+  });
+
+  // Figure out the output shape.
+  const size = computeOutShape(begin, end, strides);
+  // Remove the axes based on shrinkMask.
+  const outShape = size.filter((_, axis) => shrinkAxes.indexOf(axis) === -1);
+
   const nonStrided = strides.every(v => v === 1);
   if (nonStrided) {
-    const [beginIndex, size, shrinkAxis] = getStridedSlicedInfo(
-        $x.shape, begin, end, strides, beginMask, endMask, ellipsisMask,
-        newAxisMask, shrinkAxisMask);
-    const outShape =
-        size.filter((_, index) => shrinkAxis.indexOf(index) === -1);
-    return slice($x, beginIndex, size).reshape(outShape);
+    return slice($x, begin, size).reshape(outShape);
   }
-  return ENGINE.runKernel(
-      backend => backend.stridedSlice(
-          $x, begin, end, strides, beginMask, endMask, ellipsisMask,
-          newAxisMask, shrinkAxisMask),
-      {$x});
+  const res = ENGINE.runKernel(
+      backend => backend.stridedSlice($x, begin, end, strides), {$x});
+  return res.reshape(outShape);
 }
 
 export const stridedSlice = op({stridedSlice_});

tfjs-core/src/ops/strided_slice_test.ts

@@ -20,9 +20,53 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('stridedSlice', ALL_ENVS, () => {
-  it('stridedSlice should fail if new axis mask is set', () => {
-    const tensor = tf.tensor1d([0, 1, 2, 3]);
-    expect(() => tf.stridedSlice(tensor, [0], [3], [2], 0, 0, 0, 1)).toThrow();
+  it('stridedSlice with first axis being new', async () => {
+    // Python slice code: t[tf.newaxis,0:3]
+    const t = tf.tensor1d([0, 1, 2, 3]);
+    const begin = [0, 0];
+    const end = [1, 3];
+    const strides = [1, 2];
+    const beginMask = 0;
+    const endMask = 0;
+    const ellipsisMask = 0;
+    const newAxisMask = 1;
+
+    const output = tf.stridedSlice(
+        t, begin, end, strides, beginMask, endMask, ellipsisMask, newAxisMask);
+    expect(output.shape).toEqual([1, 2]);
+    expectArraysClose(await output.data(), [0, 2]);
+  });
+
+  it('strided slice with several new axes', () => {
+    // Python slice code: t[1:2,tf.newaxis,0:3,tf.newaxis,2:5]
+    const t = tf.zeros([2, 3, 4, 5]);
+    const begin = [1, 0, 0, 0, 2];
+    const end = [2, 1, 3, 1, 5];
+    const strides: number[] = null;
+    const beginMask = 0;
+    const endMask = 0;
+    const ellipsisMask = 0;
+    const newAxisMask = 0b1010;
+    const output = tf.stridedSlice(
+        t, begin, end, strides, beginMask, endMask, ellipsisMask, newAxisMask);
+    expect(output.shape).toEqual([1, 1, 3, 1, 2, 5]);
+  });
+
+  it('strided slice with new axes and shrink axes', () => {
+    // Python slice code: t[1:2,tf.newaxis,1,tf.newaxis,2,2:5]
+    const t = tf.zeros([2, 3, 4, 5]);
+    const begin = [1, 0, 1, 0, 2, 2];
+    const end = [2, 1, 2, 1, 3, 5];
+    const strides: number[] = null;
+    const beginMask = 0;
+    const endMask = 0;
+    const ellipsisMask = 0;
+    const newAxisMask = 0b1010;
+    const shrinkAxisMask = 0b10100;
+    const output = tf.stridedSlice(
+        t, begin, end, strides, beginMask, endMask, ellipsisMask, newAxisMask,
+        shrinkAxisMask);
+    expect(output.shape).toEqual([1, 1, 1, 3]);
   });
 
   it('stridedSlice should fail if ellipsis mask is set', () => {

tfjs-core/src/tensor_test.ts

@@ -1461,6 +1461,14 @@ describeWithFlags('tensor', ALL_ENVS, () => {
     expect(res.shape).toEqual([0, 0]);
   });
 
+  it('squeeze can take an empty list of axis', () => {
+    const a = tf.zeros([2, 1, 3, 1, 4]);
+    const axes: number[] = [];
+    // Empty axes list means all possible axes.
+    const res = tf.squeeze(a, axes);
+    expect(res.shape).toEqual([2, 3, 4]);
+  });
+
   it('squeeze a complex64 tensor', async () => {
     const a = tf.complex([[4], [1], [5]], [[2], [3], [6]]);
     const b = a.squeeze();

tfjs-core/src/util.ts

@@ -349,7 +349,10 @@ export function squeezeShape(shape: number[], axis?: number[]):
     {newShape: number[], keptDims: number[]} {
   const newShape: number[] = [];
   const keptDims: number[] = [];
-  const axes = axis == null ? null : parseAxisParam(axis, shape).sort();
+  const isEmptyArray = axis != null && Array.isArray(axis) && axis.length === 0;
+  const axes = (axis == null || isEmptyArray) ?
+      null :
+      parseAxisParam(axis, shape).sort();
   let j = 0;
   for (let i = 0; i < shape.length; ++i) {
     if (axes != null) {