Export types used in tfjs-data (tensorflow#1626)

BUG
* save

* save

Author: Kangyi Zhang

src/index.ts

@@ -30,6 +30,7 @@ import * as math from './math';
 import * as browser from './ops/browser';
 import * as serialization from './serialization';
 import {setOpHandler} from './tensor';
+import * as tensor_util from './tensor_util';
 import * as test_util from './test_util';
 import * as util from './util';
 import {version} from './version';
@@ -46,8 +47,8 @@ export {Optimizer} from './optimizers/optimizer';
 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 {GradSaveFunc, NamedTensorMap} from './tensor_types';
-export {DataType, DataTypeMap, DataValues, Rank, ShapeMap} from './types';
+export {GradSaveFunc, NamedTensorMap, TensorContainer, TensorContainerArray, TensorContainerObject} from './tensor_types';
+export {DataType, DataTypeMap, DataValues, Rank, ShapeMap, TensorLike} from './types';
 
 export * from './ops/ops';
 export {LSTMCellFunc} from './ops/lstm';
@@ -76,7 +77,17 @@ export {
 };
 
 // Second level exports.
-export {browser, environment, io, math, serialization, test_util, util, webgl};
+export {
+  browser,
+  environment,
+  io,
+  math,
+  serialization,
+  test_util,
+  util,
+  webgl,
+  tensor_util
+};
 
 // Backend specific.
 export {KernelBackend, BackendTimingInfo, DataMover, DataStorage} from './kernels/backend';

src/tensor_util.ts

@@ -16,7 +16,7 @@
  */
 
 import {Tensor} from './tensor';
-import {NamedTensorMap, TensorContainer, TensorContainerArray} from './tensor_types';
+import {TensorContainer, TensorContainerArray} from './tensor_types';
 import {upcastType} from './types';
 import {assert} from './util';
 
@@ -44,33 +44,6 @@ export function isTensorInList(tensor: Tensor, tensorList: Tensor[]): boolean {
   return false;
 }
 
-export function flattenNameArrayMap(
-    nameArrayMap: Tensor|NamedTensorMap, keys?: string[]): Tensor[] {
-  const xs: Tensor[] = [];
-  if (nameArrayMap instanceof Tensor) {
-    xs.push(nameArrayMap);
-  } else {
-    const xMap = nameArrayMap as {[xName: string]: Tensor};
-    for (let i = 0; i < keys.length; i++) {
-      xs.push(xMap[keys[i]]);
-    }
-  }
-  return xs;
-}
-
-export function unflattenToNameArrayMap(
-    keys: string[], flatArrays: Tensor[]): NamedTensorMap {
-  if (keys.length !== flatArrays.length) {
-    throw new Error(
-        `Cannot unflatten Tensor[], keys and arrays are not of same length.`);
-  }
-  const result: NamedTensorMap = {};
-  for (let i = 0; i < keys.length; i++) {
-    result[keys[i]] = flatArrays[i];
-  }
-  return result;
-}
-
 /**
  * Extracts any `Tensor`s found within the provided object.
  *

src/tensor_util_test.ts

@@ -18,8 +18,7 @@
 import * as tf from './index';
 import {describeWithFlags} from './jasmine_util';
 import {Tensor} from './tensor';
-import {NamedTensorMap} from './tensor_types';
-import {flattenNameArrayMap, getTensorsInContainer, isTensorInList, unflattenToNameArrayMap} from './tensor_util';
+import {getTensorsInContainer, isTensorInList} from './tensor_util';
 import {convertToTensor} from './tensor_util_env';
 import {ALL_ENVS, expectArraysClose, expectArraysEqual} from './test_util';
 
@@ -39,29 +38,6 @@ describe('tensor_util.isTensorInList', () => {
   });
 });
 
-describe('tensor_util.flattenNameArrayMap', () => {
-  it('basic', () => {
-    const a = tf.scalar(1);
-    const b = tf.scalar(3);
-    const c = tf.tensor1d([1, 2, 3]);
-
-    const map: NamedTensorMap = {a, b, c};
-    expect(flattenNameArrayMap(map, Object.keys(map))).toEqual([a, b, c]);
-  });
-});
-
-describe('tensor_util.unflattenToNameArrayMap', () => {
-  it('basic', () => {
-    const a = tf.scalar(1);
-    const b = tf.scalar(3);
-    const c = tf.tensor1d([1, 2, 3]);
-
-    expect(unflattenToNameArrayMap(['a', 'b', 'c'], [
-      a, b, c
-    ])).toEqual({a, b, c});
-  });
-});
-
 describe('getTensorsInContainer', () => {
   it('null input returns empty tensor', () => {
     const results = getTensorsInContainer(null);