Add string dtype to Tensor

.vscode/launch.json

@@ -0,0 +1,19 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "type": "chrome",
+      "request": "attach",
+      "name": "Attach Karma Chrome",
+      "address": "localhost",
+      "port": 9333,
+      "pathMapping": {
+        "/": "${workspaceRoot}",
+        "/base/": "${workspaceRoot}/"
+      }
+    }
+  ]
+}

karma.conf.js

@@ -89,7 +89,9 @@ module.exports = function(config) {
       chrome_with_swift_shader: {
         base: 'Chrome',
         flags: ['--blacklist-accelerated-compositing', '--blacklist-webgl']
-      }
+      },
+      chrome_debugging:
+          {base: 'Chrome', flags: ['--remote-debugging-port=9333']}
     },
     client: {jasmine: {random: false}, args: args}
   });

src/buffer_test.ts

@@ -0,0 +1,73 @@
+/**
+ * @license
+ * Copyright 2018 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tf from './index';
+import {describeWithFlags} from './jasmine_util';
+import {ALL_ENVS, expectArraysClose, expectArraysEqual} from './test_util';
+
+describeWithFlags('tf.buffer', ALL_ENVS, () => {
+  it('float32', () => {
+    const buff = tf.buffer([1, 2, 3], 'float32');
+    buff.set(1, 0, 0, 0);
+    buff.set(2, 0, 1, 0);
+    expect(buff.get(0, 0, 0)).toEqual(1);
+    expect(buff.get(0, 0, 1)).toEqual(0);
+    expect(buff.get(0, 0, 2)).toEqual(0);
+    expect(buff.get(0, 1, 0)).toEqual(2);
+    expect(buff.get(0, 1, 1)).toEqual(0);
+    expect(buff.get(0, 1, 2)).toEqual(0);
+    expectArraysClose(buff.toTensor(), [1, 0, 0, 2, 0, 0]);
+    expectArraysClose(buff.values, new Float32Array([1, 0, 0, 2, 0, 0]));
+  });
+
+  it('int32', () => {
+    const buff = tf.buffer([2, 3], 'int32');
+    buff.set(1.3, 0, 0);
+    buff.set(2.1, 1, 1);
+    expect(buff.get(0, 0)).toEqual(1);
+    expect(buff.get(0, 1)).toEqual(0);
+    expect(buff.get(0, 2)).toEqual(0);
+    expect(buff.get(1, 0)).toEqual(0);
+    expect(buff.get(1, 1)).toEqual(2);
+    expect(buff.get(1, 2)).toEqual(0);
+    expectArraysClose(buff.toTensor(), [1, 0, 0, 0, 2, 0]);
+    expectArraysClose(buff.values, new Int32Array([1, 0, 0, 0, 2, 0]));
+  });
+
+  it('bool', () => {
+    const buff = tf.buffer([4], 'bool');
+    buff.set(true, 1);
+    buff.set(true, 2);
+    expect(buff.get(0)).toBeFalsy();
+    expect(buff.get(1)).toBeTruthy();
+    expect(buff.get(2)).toBeTruthy();
+    expect(buff.get(3)).toBeFalsy();
+    expectArraysClose(buff.toTensor(), [0, 1, 1, 0]);
+    expectArraysClose(buff.values, new Uint8Array([0, 1, 1, 0]));
+  });
+
+  it('string', () => {
+    const buff = tf.buffer([2, 2], 'string');
+    buff.set('first', 0, 0);
+    buff.set('third', 1, 0);
+    expect(buff.get(0, 0)).toEqual('first');
+    expect(buff.get(0, 1)).toBeFalsy();
+    expect(buff.get(1, 0)).toEqual('third');
+    expect(buff.get(1, 1)).toBeFalsy();
+    expectArraysEqual(buff.toTensor(), ['first', null, 'third', null]);
+  });
+});

src/engine.ts

@@ -21,9 +21,9 @@ import {backpropagateGradients, getFilteredNodesXToY, NamedGradientMap, TapeNode
 import {DataId, Tensor, Tensor3D, Variable} from './tensor';
 import {NamedTensorMap, NamedVariableMap, TensorContainer} from './tensor_types';
 import {getTensorsInContainer, isTensorInList} from './tensor_util';
-import {DataType, TypedArray} from './types';
+import {DataType, DataValues} from './types';
 import * as util from './util';
-import {makeOnesTypedArray, now, sizeFromShape} from './util';
+import {bytesFromStringArray, makeOnesTypedArray, now, sizeFromShape} from './util';
 
 /**
  * A function that computes an output. The save function is for saving tensors
@@ -253,11 +253,11 @@ export class Engine implements TensorManager, DataMover {
     if (refCount === 0) {
       this.numDataBuffers++;
 
-      // Don't count bytes for complex numbers as they are counted by their
-      // components.
-      if (a.dtype !== 'complex64') {
-        this.numBytes +=
-            util.sizeFromShape(a.shape) * util.bytesPerElement(a.dtype);
+      // Bytes for complex numbers are counted by their
+      // components. Bytes for string tensors are counted when someone
+      // writes values.
+      if (a.dtype !== 'complex64' && a.dtype !== 'string') {
+        this.numBytes += a.size * util.bytesPerElement(a.dtype);
       }
       this.tensorInfo.set(
           a.dataId,
@@ -287,15 +287,17 @@ export class Engine implements TensorManager, DataMover {
     this.numTensors--;
     const refCount = this.tensorInfo.get(a.dataId).refCount;
     if (refCount <= 1) {
-      const info = this.tensorInfo.get(a.dataId);
-      info.backend.disposeData(a.dataId);
-      this.numDataBuffers--;
-      // Don't count bytes for complex numbers as they are counted by their
-      // components.
-      if (a.dtype !== 'complex64') {
+      if (a.dtype === 'string') {
+        this.numBytes -= bytesFromStringArray(a.dataSync<'string'>());
+      } else if (a.dtype !== 'complex64') {
+        // Don't count bytes for complex numbers as they are counted by their
+        // components.
         this.numBytes -=
             util.sizeFromShape(a.shape) * util.bytesPerElement(a.dtype);
       }
+      this.numDataBuffers--;
+      const info = this.tensorInfo.get(a.dataId);
+      info.backend.disposeData(a.dataId);
       this.tensorInfo.delete(a.dataId);
     } else {
       this.tensorInfo.get(a.dataId).refCount--;
@@ -537,8 +539,14 @@ export class Engine implements TensorManager, DataMover {
   }
 
   // Forwarding to backend.
-  write(dataId: DataId, values: TypedArray): void {
+  write(dataId: DataId, values: DataValues): void {
     const info = this.tensorInfo.get(dataId);
+    if (info.dtype === 'string') {
+      const oldBytes =
+          bytesFromStringArray(info.backend.readSync(dataId) as string[]);
+      const newBytes = bytesFromStringArray(values as string[]);
+      this.numBytes += newBytes - oldBytes;
+    }
     if (this.backend !== info.backend) {
       // Delete the tensor from the old backend and move it to the new backend.
       info.backend.disposeData(dataId);
@@ -547,12 +555,12 @@ export class Engine implements TensorManager, DataMover {
     }
     this.backend.write(dataId, values);
   }
-  readSync(dataId: DataId): TypedArray {
+  readSync(dataId: DataId): DataValues {
     // Route the read to the correct backend.
     const info = this.tensorInfo.get(dataId);
     return info.backend.readSync(dataId);
   }
-  read(dataId: DataId): Promise<TypedArray> {
+  read(dataId: DataId): Promise<DataValues> {
     // Route the read to the correct backend.
     const info = this.tensorInfo.get(dataId);
     return info.backend.read(dataId);

src/engine_test.ts

@@ -361,6 +361,18 @@ describeWithFlags('memory', ALL_ENVS, () => {
     expect(sum.dtype).toBe('int32');
     expectArraysClose(sum, [1 + 1 + 0 + 1]);
   });
+
+  it('string tensor', () => {
+    const a = tf.tensor([['a', 'bb'], ['c', 'd']]);
+
+    expect(tf.memory().numTensors).toBe(1);
+    expect(tf.memory().numBytes).toBe(10);  // 10 letters, each 2 bytes.
+
+    a.dispose();
+
+    expect(tf.memory().numTensors).toBe(0);
+    expect(tf.memory().numBytes).toBe(0);
+  });
 });
 
 describeWithFlags('profile', ALL_ENVS, () => {

src/index.ts

@@ -48,7 +48,7 @@ export {RMSPropOptimizer} from './optimizers/rmsprop_optimizer';
 export {SGDOptimizer} from './optimizers/sgd_optimizer';
 export {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, TensorBuffer, variable, Variable} from './tensor';
 export {NamedTensorMap} from './tensor_types';
-export {DataType, Rank, ShapeMap} from './types';
+export {DataType, DataTypeMap, DataValues, Rank, ShapeMap} from './types';
 
 export * from './ops/ops';
 export {LSTMCellFunc} from './ops/lstm';

src/jasmine_util.ts

@@ -93,7 +93,7 @@ export interface TestEnv {
 
 export let TEST_ENVS: TestEnv[] = [
   {
-    name: 'test-webgl1',
+    name: 'webgl1',
     factory: () => new MathBackendWebGL(),
     features: {
       'WEBGL_VERSION': 1,
@@ -102,7 +102,7 @@ export let TEST_ENVS: TestEnv[] = [
     }
   },
   {
-    name: 'test-webgl2',
+    name: 'webgl2',
     factory: () => new MathBackendWebGL(),
     features: {
       'WEBGL_VERSION': 2,
@@ -111,7 +111,7 @@ export let TEST_ENVS: TestEnv[] = [
     }
   },
   {
-    name: 'test-cpu',
+    name: 'cpu',
     factory: () => new MathBackendCPU(),
     features: {'HAS_WEBGL': false}
   }

src/kernels/backend.ts

@@ -17,7 +17,7 @@
 
 import {Conv2DInfo} from '../ops/conv_util';
 import {DataId, Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor';
-import {DataType, Rank, ShapeMap, TypedArray} from '../types';
+import {DataType, DataValues, Rank, ShapeMap} from '../types';
 
 // Required information for all backends.
 export interface BackendTimingInfo {
@@ -27,10 +27,10 @@ export interface BackendTimingInfo {
 }
 
 export interface TensorStorage {
-  read(dataId: DataId): Promise<TypedArray>;
-  readSync(dataId: DataId): TypedArray;
+  read(dataId: DataId): Promise<DataValues>;
+  readSync(dataId: DataId): DataValues;
   disposeData(dataId: DataId): void;
-  write(dataId: DataId, values: TypedArray): void;
+  write(dataId: DataId, values: DataValues): void;
   fromPixels(
       pixels: ImageData|HTMLImageElement|HTMLCanvasElement|HTMLVideoElement,
       numChannels: number): Tensor3D;
@@ -87,26 +87,24 @@ export class KernelBackend implements TensorStorage, BackendTimer {
   time(f: () => void): Promise<BackendTimingInfo> {
     throw new Error('Not yet implemented.');
   }
-  read(dataId: object): Promise<Float32Array|Int32Array|Uint8Array> {
+  read(dataId: object): Promise<DataValues> {
     throw new Error('Not yet implemented.');
   }
-  readSync(dataId: object): Float32Array|Int32Array|Uint8Array {
+  readSync(dataId: object): DataValues {
     throw new Error('Not yet implemented.');
   }
   disposeData(dataId: object): void {
     throw new Error('Not yet implemented.');
   }
-  write(dataId: object, values: Float32Array|Int32Array|Uint8Array): void {
+  write(dataId: object, values: DataValues): void {
     throw new Error('Not yet implemented.');
   }
   fromPixels(
       pixels: ImageData|HTMLImageElement|HTMLCanvasElement|HTMLVideoElement,
       numChannels: number): Tensor<Rank.R3> {
     throw new Error('Not yet implemented.');
   }
-  register(
-      dataId: object, shape: number[],
-      dtype: 'float32'|'int32'|'bool'|'complex64'): void {
+  register(dataId: object, shape: number[], dtype: DataType): void {
     throw new Error('Not yet implemented.');
   }
   memory(): {unreliable: boolean;} {