Add dl.concat/stack/expandDims/pad and tensor.buffer() (#708)

- Align `dl.concat` to match `tf.concat` (takes a list of tensors, instead of just two tensors)
- Add `dl.stack` (and its chain sibling)
- Add `dl.expandDims` (and its chain sibling)
- Add `dl.pad` (and its chain sibling). Also have `dl.pad1d/2d` call into `dl.pad`.
- Deprecate `tensor.val()/get()/locToIndex()/indexToVal()`
- Add `tensor.buffer()` which returns `TensorBuffer` holding the underlying data.

Fixes #696, #706 and #705

Author: dsmilkov

karma.conf.js

@@ -34,6 +34,7 @@ module.exports = function(config) {
       username: process.env.BROWSERSTACK_USERNAME,
       accessKey: process.env.BROWSERSTACK_KEY
     },
+    reportSlowerThan: 500,
     browserNoActivityTimeout: 30000,
     customLaunchers: {
       bs_chrome_mac: {

src/math.ts

@@ -21,7 +21,6 @@ import * as array_ops from './ops/array_ops';
 import * as batchnorm from './ops/batchnorm';
 import * as binary_ops from './ops/binary_ops';
 import * as compare from './ops/compare';
-import * as concat from './ops/concat';
 import * as conv from './ops/conv';
 import * as image_ops from './ops/image_ops';
 import * as logical from './ops/logical_ops';
@@ -38,7 +37,7 @@ import * as softmax_ops from './ops/softmax';
 import * as transpose from './ops/transpose';
 import * as unary_ops from './ops/unary_ops';
 import {ScopeResult} from './tape_util';
-import {Scalar, Tensor, Tensor1D} from './tensor';
+import {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from './tensor';
 import {Tracking} from './tracking';
 import {Rank} from './types';
 import * as util from './util';
@@ -66,12 +65,6 @@ export class NDArrayMath {
   reverse3D = reverse.Ops.reverse3d;
   reverse4D = reverse.Ops.reverse4d;
 
-  concat = concat.Ops.concat;
-  concat1D = concat.Ops.concat1d;
-  concat2D = concat.Ops.concat2d;
-  concat3D = concat.Ops.concat3d;
-  concat4D = concat.Ops.concat4d;
-
   batchNormalization = batchnorm.Ops.batchNormalization;
   batchNormalization2D = batchnorm.Ops.batchNormalization2d;
   batchNormalization3D = batchnorm.Ops.batchNormalization3d;
@@ -329,6 +322,31 @@ export class NDArrayMath {
             `got rank ${c.rank}.`);
     return this.multiply(c, a) as T;
   }
+
+  /** @deprecated */
+  concat<T extends Tensor>(a: T, b: T, axis: number): T {
+    return ops.concat([a, b], axis);
+  }
+
+  /** @deprecated */
+  concat1D(a: Tensor1D, b: Tensor1D): Tensor1D {
+    return ops.concat1d([a, b]);
+  }
+
+  /** @deprecated */
+  concat2D(a: Tensor2D, b: Tensor2D, axis: number): Tensor2D {
+    return ops.concat2d([a, b], axis);
+  }
+
+  /** @deprecated */
+  concat3D(a: Tensor3D, b: Tensor3D, axis: number): Tensor3D {
+    return ops.concat3d([a, b], axis);
+  }
+
+  /** @deprecated */
+  concat4D(a: Tensor4D, b: Tensor4D, axis: number): Tensor4D {
+    return ops.concat4d([a, b], axis);
+  }
 }
 
 export type ScopeFn<T extends ScopeResult> =

src/ops/array_ops.ts

@@ -20,9 +20,9 @@ import {ENV} from '../environment';
 // tslint:disable-next-line:max-line-length
 import {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, TensorBuffer} from '../tensor';
 // tslint:disable-next-line:max-line-length
-import {ArrayData, DataType, DataTypeMap, Rank, ShapeMap, TensorLike, TensorLike1D, TensorLike2D, TensorLike3D, TensorLike4D} from '../types';
+import {ArrayData, DataType, DataTypeMap, Rank, ShapeMap, TensorLike, TensorLike1D, TensorLike2D, TensorLike3D, TensorLike4D, TypedArray} from '../types';
 import * as util from '../util';
-
+import {Concat} from './concat';
 import {operation} from './operation';
 import {MPRandGauss} from './rand';
 
@@ -595,17 +595,16 @@ export class Ops {
   }
 
   /**
-   * Pads a Tensor1D with a given value.
+   * Pads a `Tensor1D` with a given value.
    *
    * This operation will pad a tensor according to the `paddings` you specify.
    *
    * This operation currently only implements the `CONSTANT` mode from
    * Tensorflow's `pad` operation.
    *
    * @param x The tensor to pad.
-   * @param paddings A tuple of ints [padLeft, padRight], how much to pad on the
-   *     left and right side of the tensor.
-   * @param constantValue The scalar pad value to use. Defaults to 0.
+   * @param paddings A tuple of ints `[padLeft, padRight]`, how much to pad.
+   * @param constantValue The pad value to use. Defaults to 0.
    */
   @operation
   static pad1d(x: Tensor1D, paddings: [number, number], constantValue = 0):
@@ -618,18 +617,15 @@ export class Ops {
   }
 
   /**
-   * Pads a Tensor2D with a given value.
-   *
-   * This operation will pad a tensor according to the `paddings` you specify.
+   * Pads a `Tensor2D` with a given value and the `paddings` you specify.
    *
    * This operation currently only implements the `CONSTANT` mode from
-   * Tensorflow's `pad` operation.
+   * TensorFlow's `pad` operation.
    *
    * @param x The tensor to pad.
-   * @param paddings A pair of tuple ints
-   *     [[padTop, padBottom], [padLeft, padRight]], how much to pad on the
-   *     tensor.
-   * @param constantValue The scalar pad value to use. Defaults to 0.
+   * @param paddings A pair of tuple ints:
+   *     `[[padTop, padBottom], [padLeft, padRight]]`, how much to pad.
+   * @param constantValue The pad value to use. Defaults to 0.
    */
   @operation
   static pad2d(
@@ -643,6 +639,83 @@ export class Ops {
         'Pad2D', {inputs: {x}, args: {paddings, constantValue}});
   }
 
+  /**
+   * Pads a `Tensor` with a given value and the `paddings` you specify.
+   *
+   * This operation currently only implements the `CONSTANT` mode from
+   * Tensorflow's `pad` operation.
+   *
+   * @param x The tensor to pad.
+   * @param paddings An array of length `R` (the rank of the tensor), where each
+   *     element is a length-2 tuple of ints `[padBefore, padAfter]`, specifying
+   *     how much to pad along each dimension of the tensor.
+   * @param constantValue The pad value to use. Defaults to 0.
+   */
+  @doc({heading: 'Tensors', subheading: 'Transformations'})
+  @operation
+  static pad<T extends Tensor>(
+      x: T, paddings: Array<[number, number]>, constantValue = 0): T {
+    if (x.rank === 0) {
+      throw new Error('pad(scalar) is not defined. Pass non-scalar to pad');
+    } else if (x.rank === 1) {
+      return Ops.pad1d(x as Tensor1D, paddings[0], constantValue) as T;
+    } else if (x.rank === 2) {
+      return Ops.pad2d(
+                 x as Tensor2D,
+                 paddings as [[number, number], [number, number]],
+                 constantValue) as T;
+    } else {
+      throw new Error(`pad of rank-${x.rank} tensor is not yet supported`);
+    }
+  }
+
+  /**
+   * Stacks a list of rank-`R` `Tensor`s into one rank-`(R+1)` `Tensor`.
+   *
+   * @param tensors A list of tensor objects with the same shape and dtype.
+   * @param axis The axis to stack along. Defaults to 0 (the first dim).
+   */
+  @doc({heading: 'Tensors', subheading: 'Transformations'})
+  @operation
+  static stack<T extends Tensor>(tensors: T[], axis = 0): Tensor {
+    util.assert(tensors.length >= 2, 'Pass at least two tensors to dl.stack');
+    const rank = tensors[0].rank;
+    const shape = tensors[0].shape;
+    const dtype = tensors[0].dtype;
+
+    util.assert(axis <= rank, 'Axis must be <= rank of the tensor');
+
+    tensors.forEach(t => {
+      util.assertShapesMatch(
+          shape, t.shape,
+          'All tensors passed to stack must have matching shapes');
+    });
+
+    tensors.forEach(t => {
+      util.assert(
+          dtype === t.dtype,
+          'All tensors passed to stack must have matching dtypes');
+    });
+    const expandedTensors = tensors.map(t => t.expandDims(axis));
+    return Concat.concat(expandedTensors, axis);
+  }
+
+  /**
+   * Returns a `Tensor` that has expanded rank, by inserting a dimension
+   * into the tensor's shape.
+   *
+   * @param axis The dimension index at which to insert shape of `1`. Defaults
+   *     to 0 (the first dimension).
+   */
+  @doc({heading: 'Tensors', subheading: 'Transformations'})
+  @operation
+  static expandDims<R2 extends Rank>(x: Tensor, axis = 0): Tensor<R2> {
+    util.assert(axis <= x.rank, 'Axis must be <= rank of the tensor');
+    const newShape = x.shape.slice();
+    newShape.splice(axis, 0, 1);
+    return Ops.reshape(x, newShape);
+  }
+
   /**
    * Return an evenly spaced sequence of numbers over the given interval.
    *
@@ -720,18 +793,19 @@ export class Ops {
   /**
    * Creates an empty `TensorBuffer` with the specified `shape` and `dtype`.
    *
-   * The values are stored in cpu as a `TypedArray`. Fill the buffer using
-   * `buffer.set()`, or by modifying directly `buffer.values`.
+   * The values are stored in cpu as `TypedArray`. Fill the buffer using
+   * `buffer.set()`, or by modifying directly `buffer.values`. When done,
+   * call `buffer.toTensor()` to get an immutable `Tensor` with those values.
    *
-   * When done, call `buffer.toTensor()` to get an immutable `Tensor` with those
-   * values.
    * @param shape An array of integers defining the output tensor shape.
    * @param dtype The dtype of the buffer. Defaults to 'float32'.
+   * @param values The values of the buffer as `TypedArray`. Defaults to zeros.
    */
   @doc({heading: 'Tensors', subheading: 'Creation'})
   static buffer<R extends Rank>(
-      shape: ShapeMap[R], dtype: DataType = 'float32'): TensorBuffer<R> {
-    return new TensorBuffer<R>(shape, dtype);
+      shape: ShapeMap[R], dtype: DataType = 'float32', values?: TypedArray):
+      TensorBuffer<R> {
+    return new TensorBuffer<R>(shape, dtype, values);
   }
 
   /**

src/ops/array_ops_test.ts

@@ -17,10 +17,9 @@
 
 import * as dl from '../index';
 // tslint:disable-next-line:max-line-length
-import {expectArrayInMeanStdRange, jarqueBeraNormalityTest} from './rand_util';
-// tslint:disable-next-line:max-line-length
 import {ALL_ENVS, describeWithFlags, expectArraysClose, expectArraysEqual, expectValuesInRange} from '../test_util';
 import * as util from '../util';
+import {expectArrayInMeanStdRange, jarqueBeraNormalityTest} from './rand_util';
 
 describeWithFlags('zeros', ALL_ENVS, () => {
   it('1D default dtype', () => {
@@ -1632,3 +1631,121 @@ describeWithFlags('fill', ALL_ENVS, () => {
     expectArraysClose(a, [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]);
   });
 });
+
+describeWithFlags('stack', ALL_ENVS, () => {
+  it('scalars 3, 5 and 7', () => {
+    const a = dl.scalar(3);
+    const b = dl.scalar(5);
+    const c = dl.scalar(7);
+    const res = dl.stack([a, b, c]);
+    expect(res.shape).toEqual([3]);
+    expectArraysClose(res, [3, 5, 7]);
+  });
+
+  it('scalars 3, 5 and 7 along axis=1 throws error', () => {
+    const a = dl.scalar(3);
+    const b = dl.scalar(5);
+    const c = dl.scalar(7);
+    const f = () => dl.stack([a, b, c], 1);
+    expect(f).toThrowError();
+  });
+
+  it('non matching shapes throws error', () => {
+    const a = dl.scalar(3);
+    const b = dl.tensor1d([5]);
+    const f = () => dl.stack([a, b]);
+    expect(f).toThrowError();
+  });
+
+  it('non matching dtypes throws error', () => {
+    const a = dl.scalar(3);
+    const b = dl.scalar(5, 'bool');
+    const f = () => dl.stack([a, b]);
+    expect(f).toThrowError();
+  });
+
+  it('2d but axis=3 throws error', () => {
+    const a = dl.zeros([2, 2]);
+    const b = dl.zeros([2, 2]);
+    const f = () => dl.stack([a, b], 3 /* axis */);
+    expect(f).toThrowError();
+  });
+
+  it('[1,2], [3,4] and [5,6], axis=0', () => {
+    const a = dl.tensor1d([1, 2]);
+    const b = dl.tensor1d([3, 4]);
+    const c = dl.tensor1d([5, 6]);
+    const res = dl.stack([a, b, c], 0 /* axis */);
+    expect(res.shape).toEqual([3, 2]);
+    expectArraysClose(res, [1, 2, 3, 4, 5, 6]);
+  });
+
+  it('[1,2], [3,4] and [5,6], axis=1', () => {
+    const a = dl.tensor1d([1, 2]);
+    const b = dl.tensor1d([3, 4]);
+    const c = dl.tensor1d([5, 6]);
+    const res = dl.stack([a, b, c], 1 /* axis */);
+    expect(res.shape).toEqual([2, 3]);
+    expectArraysClose(res, [1, 3, 5, 2, 4, 6]);
+  });
+
+  it('[[1,2],[3,4]] and [[5, 6], [7, 8]], axis=0', () => {
+    const a = dl.tensor2d([[1, 2], [3, 4]]);
+    const b = dl.tensor2d([[5, 6], [7, 8]]);
+    const res = dl.stack([a, b], 0 /* axis */);
+    expect(res.shape).toEqual([2, 2, 2]);
+    expectArraysClose(res, [1, 2, 3, 4, 5, 6, 7, 8]);
+  });
+
+  it('[[1,2],[3,4]] and [[5, 6], [7, 8]], axis=2', () => {
+    const a = dl.tensor2d([[1, 2], [3, 4]]);
+    const b = dl.tensor2d([[5, 6], [7, 8]]);
+    const c = dl.tensor2d([[9, 10], [11, 12]]);
+    const res = dl.stack([a, b, c], 2 /* axis */);
+    expect(res.shape).toEqual([2, 2, 3]);
+    expectArraysClose(res, [1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12]);
+  });
+});
+
+describeWithFlags('expandDims', ALL_ENVS, () => {
+  it('scalar, default axis is 0', () => {
+    const res = dl.scalar(1).expandDims();
+    expect(res.shape).toEqual([1]);
+    expectArraysClose(res, [1]);
+  });
+
+  it('scalar, axis is out of bounds throws error', () => {
+    const f = () => dl.scalar(1).expandDims(1);
+    expect(f).toThrowError();
+  });
+
+  it('1d, axis=0', () => {
+    const res = dl.tensor1d([1, 2, 3]).expandDims(0 /* axis */);
+    expect(res.shape).toEqual([1, 3]);
+    expectArraysClose(res, [1, 2, 3]);
+  });
+
+  it('1d, axis=1', () => {
+    const res = dl.tensor1d([1, 2, 3]).expandDims(1 /* axis */);
+    expect(res.shape).toEqual([3, 1]);
+    expectArraysClose(res, [1, 2, 3]);
+  });
+
+  it('2d, axis=0', () => {
+    const res = dl.tensor2d([[1, 2], [3, 4], [5, 6]]).expandDims(0 /* axis */);
+    expect(res.shape).toEqual([1, 3, 2]);
+    expectArraysClose(res, [1, 2, 3, 4, 5, 6]);
+  });
+
+  it('2d, axis=1', () => {
+    const res = dl.tensor2d([[1, 2], [3, 4], [5, 6]]).expandDims(1 /* axis */);
+    expect(res.shape).toEqual([3, 1, 2]);
+    expectArraysClose(res, [1, 2, 3, 4, 5, 6]);
+  });
+
+  it('2d, axis=2', () => {
+    const res = dl.tensor2d([[1, 2], [3, 4], [5, 6]]).expandDims(2 /* axis */);
+    expect(res.shape).toEqual([3, 2, 1]);
+    expectArraysClose(res, [1, 2, 3, 4, 5, 6]);
+  });
+});