Allow users to pass a list of TypedArray in tf.tensor() (tensorflow#1424)

Fixes tensorflow/tfjs#959

FEATURE

Author: dsmilkov

src/kernels/backend_cpu_test.ts

@@ -64,7 +64,6 @@ describe('backendCPU', () => {
   });
 
   it('register string tensor with values and mismatched shape', () => {
-    expect(() => tf.Tensor.make([4], {values: ['a', 'b', 'c']}, 'string'))
-        .toThrowError();
+    expect(() => tf.tensor(['a', 'b', 'c'], [4], 'string')).toThrowError();
   });
 });

src/kernels/backend_webgl_test.ts

@@ -63,11 +63,12 @@ describeWithFlags('lazy packing and unpacking', WEBGL_ENVS, () => {
   });
 
   it('should work when the same input must be represented by' +
-    'different textures', () => {
-    const a = tf.tensor1d([1, 2]);
-    const res = tf.dot(a, a);
-    expectArraysClose(res, [5]);
-  });
+         'different textures',
+     () => {
+       const a = tf.tensor1d([1, 2]);
+       const res = tf.dot(a, a);
+       expectArraysClose(res, [5]);
+     });
 });
 
 describeWithFlags('backendWebGL', WEBGL_ENVS, () => {
@@ -124,9 +125,7 @@ describeWithFlags('backendWebGL', WEBGL_ENVS, () => {
     const backend = new MathBackendWebGL();
     tf.ENV.registerBackend('test-storage', () => backend);
     tf.setBackend('test-storage');
-
-    expect(() => tf.Tensor.make([4], {values: ['a', 'b', 'c']}, 'string'))
-        .toThrowError();
+    expect(() => tf.tensor(['a', 'b', 'c'], [4], 'string')).toThrowError();
   });
 
   it('delayed storage, reading', () => {

src/ops/array_ops.ts

@@ -288,7 +288,8 @@ function multinomial_(
 function oneHot_(
     indices: Tensor1D|TensorLike1D, depth: number, onValue = 1,
     offValue = 0): Tensor2D {
-  const $indices = convertToTensor(indices, 'indices', 'oneHot', 'int32');
+  const $indices =
+      convertToTensor(indices, 'indices', 'oneHot', 'int32') as Tensor1D;
 
   if (depth < 2) {
     throw new Error(`Error in oneHot: depth must be >=2, but it is ${depth}`);
@@ -1061,7 +1062,7 @@ function expandDims_<R2 extends Rank>(
 function depthToSpace_(
     x: Tensor4D|TensorLike4D, blockSize: number,
     dataFormat: 'NHWC'|'NCHW' = 'NHWC'): Tensor4D {
-  const $x = convertToTensor(x, 'x', 'depthToSpace');
+  const $x = convertToTensor(x, 'x', 'depthToSpace') as Tensor4D;
 
   const inputHeight = (dataFormat === 'NHWC') ? $x.shape[1] : $x.shape[2];
   const inputWidth = (dataFormat === 'NHWC') ? $x.shape[2] : $x.shape[3];

src/ops/tensor_ops.ts

@@ -20,7 +20,8 @@ import {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, Tensor
 import {convertToTensor, inferShape} from '../tensor_util_env';
 import {TensorLike, TensorLike1D, TensorLike2D, TensorLike3D, TensorLike4D, TensorLike5D, TensorLike6D, TypedArray} from '../types';
 import {DataType, Rank, ShapeMap} from '../types';
-import {assertNonNull, assertShapesMatch, flatten, getArrayFromDType, inferDtype, isTypedArray, makeOnesTypedArray, makeZerosTypedArray, sizeFromShape, toTypedArray} from '../util';
+import {assert, assertNonNull, flatten, getArrayFromDType, inferDtype, isTypedArray, makeOnesTypedArray, makeZerosTypedArray, sizeFromShape, toTypedArray} from '../util';
+
 import {complex} from './complex_ops';
 import {op} from './operation';
 
@@ -68,19 +69,33 @@ function tensor<R extends Rank>(
         'an array of numbers/booleans/strings, or a TypedArray');
   }
   const inferredShape = inferShape(values);
-  if (shape != null && inferredShape.length !== 1) {
-    assertShapesMatch(
-        shape, inferredShape,
-        `Error creating a new Tensor. ` +
-            `Inferred shape (${inferredShape}) does not match the ` +
-            `provided shape (${shape}). `);
+  if (shape != null) {
+    const providedSize = sizeFromShape(shape);
+    const inferredSize = sizeFromShape(inferredShape);
+    assert(
+        providedSize === inferredSize,
+        () =>
+            `Based on the provided shape, [${shape}], the tensor should have ` +
+            `${providedSize} values but has ${inferredSize}`);
+
+    for (let i = 0; i < inferredShape.length; ++i) {
+      const inferred = inferredShape[i];
+      const flatDimsDontMatch = i === inferredShape.length - 1 ?
+          inferred !== sizeFromShape(shape.slice(i)) :
+          true;
+      assert(
+          inferredShape[i] === shape[i] || !flatDimsDontMatch,
+          () => `Error creating a new Tensor. Inferred shape ` +
+              `(${inferredShape}) does not match the provided ` +
+              `shape (${shape}). `);
+    }
   }
   if (!isTypedArray(values) && !Array.isArray(values)) {
     values = [values] as number[];
   }
   shape = shape || inferredShape;
   values = dtype !== 'string' ? toTypedArray(values, dtype, ENV.get('DEBUG')) :
-                                flatten(values) as string[];
+                                flatten(values as string[]) as string[];
   return Tensor.make(shape, {values}, dtype);
 }
 

src/tensor.ts

@@ -410,14 +410,6 @@ export class Tensor<R extends Rank = Rank> {
     this.shape = shape.slice();
     this.dtype = dtype || 'float32';
     this.size = util.sizeFromShape(shape);
-    if (values != null) {
-      util.assert(
-          this.size === values.length,
-          `Based on the provided shape, [${shape}], and dtype ` +
-              `${this.dtype}, the tensor should have ` +
-              `${this.size} values but has ${values.length}`);
-    }
-
     this.strides = computeStrides(shape);
     this.dataId = dataId != null ? dataId : {};
     this.id = trackerFn().nextTensorId();

src/tensor_test.ts

@@ -575,6 +575,64 @@ describeWithFlags('tensor', ALL_ENVS, () => {
     expectArraysClose(a, [1, 5, 2]);
   });
 
+  it('tf.tensor() from Float32Array and number[]', () => {
+    const a = tf.tensor([
+      new Float32Array([1, 2]), new Float32Array([3, 4]),
+      new Float32Array([5, 6]), [7, 8]
+    ]);
+    expect(a.dtype).toBe('float32');
+    expect(a.shape).toEqual([4, 2]);
+    expectArraysClose(a, [1, 2, 3, 4, 5, 6, 7, 8]);
+  });
+
+  it('tf.tensor() from Int32Array and number[]', () => {
+    const a = tf.tensor([
+      new Int32Array([1, 2]), new Int32Array([3, 4]), new Int32Array([5, 6]),
+      [7, 8]
+    ]);
+    expect(a.dtype).toBe('int32');
+    expect(a.shape).toEqual([4, 2]);
+    expectArraysClose(a, [1, 2, 3, 4, 5, 6, 7, 8]);
+  });
+
+  it('tf.tensor() from mixed TypedArray', () => {
+    const a = tf.tensor([
+      new Float32Array([1, 2]), new Int32Array([3, 4]), new Uint8Array([5, 6]),
+      [7, 8]
+    ]);
+    expect(a.dtype).toBe('float32');
+    expect(a.shape).toEqual([4, 2]);
+    expectArraysClose(a, [1, 2, 3, 4, 5, 6, 7, 8]);
+  });
+
+  it('tf.tensor() from TypedArrays which are themselves 3D', () => {
+    // 2 tensors, each with shape 20x20x3, as flat Float32Arrays.
+    const img1 = new Float32Array(20 * 20 * 3);
+    const img2 = new Float32Array(20 * 20 * 3);
+    const t = tf.tensor([img1, img2], [2, 20, 20, 3]);
+    expect(t.dtype).toBe('float32');
+    expect(t.shape).toEqual([2, 20, 20, 3]);
+  });
+
+  it('tf.tensor() from TypeedArrays which are themselves 3D, wrong shape',
+     () => {
+       const img1 = new Float32Array(20 * 20 * 3);
+       const img2 = new Float32Array(20 * 20 * 3);
+       expect(() => tf.tensor([img1, img2], [3, 20, 20, 3])).toThrowError();
+     });
+
+  it('tf.tensor() from TypedArray + number[] fails due to wrong shape', () => {
+    expect(() => tf.tensor([
+      new Float32Array([1, 2]),
+      new Float32Array([3, 4]),
+      new Float32Array([5, 6]),
+      // Should be of length 4
+      [7, 8, 9, 10],
+    ]))
+        .toThrowError(
+            /Element arr\[3\] should have 2 elements, but has 4 elements/);
+  });
+
   it('default dtype from ascii string', () => {
     const a = tf.tensor('hello');
     expect(a.dtype).toBe('string');
@@ -1951,15 +2009,6 @@ describeWithFlags('tensor with 0 in shape', ALL_ENVS, () => {
     expectArraysEqual(a, []);
   });
 
-  it('1d throws when values are not empty', () => {
-    const values = new Float32Array([1, 2, 3]);
-    // Have to use Tensor.make since tensor1d() does not let us provide a shape.
-    expect(() => Tensor.make([0], {values}, 'float32'))
-        .toThrowError(
-            'Based on the provided shape, [0], and dtype float32, the tensor ' +
-            'should have 0 values but has 3');
-  });
-
   it('2d of shape [0, 5]', () => {
     const a = tf.tensor2d([], [0, 5]);
     expect(a.dtype).toBe('float32');
@@ -1980,7 +2029,7 @@ describeWithFlags('tensor with 0 in shape', ALL_ENVS, () => {
     const values = [1, 2, 3, 4];
     expect(() => tf.tensor2d(values, [0, 5], 'float32'))
         .toThrowError(
-            'Based on the provided shape, [0,5], and dtype float32, the ' +
+            'Based on the provided shape, [0,5], the ' +
             'tensor should have 0 values but has 4');
   });
 
@@ -1996,7 +2045,7 @@ describeWithFlags('tensor with 0 in shape', ALL_ENVS, () => {
     const values = [1, 2, 3];
     expect(() => tf.tensor3d(values, [0, 3, 0], 'float32'))
         .toThrowError(
-            'Based on the provided shape, [0,3,0], and dtype float32, the ' +
+            'Based on the provided shape, [0,3,0], the ' +
             'tensor should have 0 values but has 3');
   });
 
@@ -2012,7 +2061,7 @@ describeWithFlags('tensor with 0 in shape', ALL_ENVS, () => {
     const values = [1, 2, 3];
     expect(() => tf.tensor4d(values, [1, 3, 0, 5], 'float32'))
         .toThrowError(
-            'Based on the provided shape, [1,3,0,5], and dtype float32, the ' +
+            'Based on the provided shape, [1,3,0,5], the ' +
             'tensor should have 0 values but has 3');
   });
 

src/tensor_util_env.ts

@@ -31,7 +31,7 @@ export function inferShape(val: TensorLike): number[] {
   }
   const shape: number[] = [];
 
-  while (firstElem instanceof Array) {
+  while (firstElem instanceof Array || isTypedArray(firstElem)) {
     shape.push(firstElem.length);
     firstElem = firstElem[0];
   }
@@ -45,11 +45,11 @@ export function inferShape(val: TensorLike): number[] {
 function deepAssertShapeConsistency(
     val: TensorLike, shape: number[], indices: number[]) {
   indices = indices || [];
-  if (!(val instanceof Array)) {
+  if (!(val instanceof Array) && !isTypedArray(val)) {
     assert(
         shape.length === 0,
         () => `Element arr[${indices.join('][')}] is a primitive, ` +
-            `but should be an array of ${shape[0]} elements`);
+            `but should be an array/TypedArray of ${shape[0]} elements`);
     return;
   }
   assert(
@@ -100,15 +100,15 @@ export function convertToTensor<T extends Tensor>(
       typeof x !== 'boolean' && typeof x !== 'string') {
     throw new Error(
         `Argument '${argName}' passed to '${functionName}' must be a ` +
-        `Tensor or TensorLike, but got '${x.constructor.name}'`);
+        `Tensor or TensorLike, but got '${(x as {}).constructor.name}'`);
   }
   const inferredShape = inferShape(x);
   if (!isTypedArray(x) && !Array.isArray(x)) {
     x = [x] as number[];
   }
   const values = inferredDtype !== 'string' ?
       toTypedArray(x, inferredDtype as DataType, ENV.get('DEBUG')) :
-      flatten(x) as string[];
+      flatten(x as string[]) as string[];
   return Tensor.make(inferredShape, {values}, inferredDtype);
 }
 

src/types.ts

@@ -120,7 +120,8 @@ export function sumOutType(type: DataType): DataType {
 }
 
 /** @docalias TypedArray|Array */
-export type TensorLike = TypedArray|number|boolean|string|RegularArray<number>|
+export type TensorLike =
+    TypedArray|number|boolean|string|RegularArray<number|number[]|TypedArray>|
     RegularArray<boolean>|RegularArray<string>;
 /** @docalias TypedArray|Array */
 export type TensorLike1D = TypedArray|number[]|boolean[]|string[];

src/util.ts

@@ -98,9 +98,10 @@ export function assertNonNull(a: TensorLike): void {
 // NOTE: We explicitly type out what T extends instead of any so that
 // util.flatten on a nested array of number doesn't try to infer T as a
 // number[][], causing us to explicitly type util.flatten<number>().
-export function flatten<T extends number|boolean|string|Promise<number>>(
+export function
+flatten<T extends number|boolean|string|Promise<number>|TypedArray>(
     arr: T|RecursiveArray<T>, ret: T[] = []): T[] {
-  if (Array.isArray(arr)) {
+  if (Array.isArray(arr) || isTypedArray(arr)) {
     for (let i = 0; i < arr.length; ++i) {
       flatten(arr[i], ret);
     }
@@ -376,7 +377,7 @@ export function hasEncodingLoss(oldType: DataType, newType: DataType): boolean {
   return true;
 }
 
-export function isTypedArray(a: TensorLike): boolean {
+export function isTypedArray(a: {}): a is Float32Array|Int32Array|Uint8Array {
   return a instanceof Float32Array || a instanceof Int32Array ||
       a instanceof Uint8Array;
 }
@@ -422,24 +423,18 @@ export function isNumber(value: {}): boolean {
 }
 
 export function inferDtype(values: TensorLike): DataType {
+  if (values instanceof Array) {
+    return inferDtype(values[0]);
+  }
   if (values instanceof Float32Array) {
     return 'float32';
   } else if (values instanceof Int32Array || values instanceof Uint8Array) {
     return 'int32';
-  } else if (
-      isNumber(values) ||
-      values instanceof Array &&
-          isNumber(getFirstElemFromNestedArray(values))) {
+  } else if (isNumber(values)) {
     return 'float32';
-  } else if (
-      isString(values) ||
-      values instanceof Array &&
-          isString(getFirstElemFromNestedArray(values))) {
+  } else if (isString(values)) {
     return 'string';
-  } else if (
-      isBoolean(values) ||
-      values instanceof Array &&
-          isBoolean(getFirstElemFromNestedArray(values))) {
+  } else if (isBoolean(values)) {
     return 'bool';
   }
   return 'float32';
@@ -549,10 +544,3 @@ export function now(): number {
         'in the browser or in Node.js');
   }
 }
-
-function getFirstElemFromNestedArray(arr: TensorLike): {} {
-  while (arr instanceof Array) {
-    arr = arr[0];
-  }
-  return arr;
-}

src/util_test.ts

@@ -103,6 +103,12 @@ describe('util.flatten', () => {
       'a', 'b', 'c', 'd', 'e'
     ]);
   });
+
+  it('mixed TypedArray and number[]', () => {
+    const data =
+        [new Float32Array([1, 2]), 3, [4, 5, new Float32Array([6, 7])]];
+    expect(util.flatten(data)).toEqual([1, 2, 3, 4, 5, 6, 7]);
+  });
 });
 
 describe('util.bytesFromStringArray', () => {