/**
 * @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 '@tensorflow/tfjs-core';
import {TensorContainerObject} from '@tensorflow/tfjs-core';
import {array, DatasetContainer} from './dataset';
import * as tfd from './index';
import {iteratorFromItems, LazyIterator} from './iterators/lazy_iterator';
import {describeAllEnvs} from './util/test_utils';

class TestObjectIterator extends LazyIterator<{}> {
  data = Array.from({length: 100}, (v, k) => k);
  currentIndex = 0;

  summary() {
    return `TestObjects`;
  }

  async next(): Promise<IteratorResult<{}>> {
    if (this.currentIndex >= 100) {
      return {value: null, done: true};
    }
    const elementNumber = this.data[this.currentIndex];
    const result = {
      'number': elementNumber,
      'numberArray': [elementNumber, elementNumber ** 2, elementNumber ** 3],
      'Tensor':
          tf.tensor1d([elementNumber, elementNumber ** 2, elementNumber ** 3]),
      'Tensor2': tf.tensor2d(
          [
            elementNumber, elementNumber ** 2, elementNumber ** 3,
            elementNumber ** 4
          ],
          [2, 2]),
      'string': `Item ${elementNumber}`
    };
    this.currentIndex++;
    return {value: result, done: false};
  }
}

export class TestDataset extends tfd.Dataset<TensorContainerObject> {
  readonly size: number;
  constructor(setSize = false) {
    super();
    if (setSize) {
      this.size = 200;
    }
  }
  async iterator(): Promise<LazyIterator<{}>> {
    return new TestObjectIterator();
  }
}

// tslint:disable-next-line:no-any
function complexifyExampleAsDict(simple: any): {} {
  const v = simple['number'];
  const w = simple['numberArray'];
  const x = simple['Tensor'];
  const y = simple['Tensor2'];
  const z = simple['string'];
  return {
    a: {v, w, q: {aa: {x, y, z}, ab: {v, w, x}}},
    b: {
      ba: {baa: y, bab: z, bac: v},
      bb: {bba: w, bbb: x, bbc: y},
      bc: {bca: z, bcb: v, bcc: w}
    },
    c: {ca: {x, y, z}, cb: {v, w, x}, cc: {y, z, v}},
  };
}

describeAllEnvs('Dataset', () => {
  it('can be concatenated', async () => {
    const a = tfd.array([{'item': 1}, {'item': 2}, {'item': 3}]);
    const b = tfd.array([{'item': 4}, {'item': 5}, {'item': 6}]);
    const result = await a.concatenate(b).toArrayForTest();
    expect(result).toEqual([
      {'item': 1}, {'item': 2}, {'item': 3}, {'item': 4}, {'item': 5},
      {'item': 6}
    ]);
  });

  it('can be created by concatenating multiple underlying datasets via ' +
         'reduce',
     async () => {
       const a = tfd.array([{'item': 1}, {'item': 2}]);
       const b = tfd.array([{'item': 3}, {'item': 4}]);
       const c = tfd.array([{'item': 5}, {'item': 6}]);
       const concatenated = [a, b, c].reduce((a, b) => a.concatenate(b));
       const result = await concatenated.toArrayForTest();
       expect(result).toEqual([
         {'item': 1}, {'item': 2}, {'item': 3}, {'item': 4}, {'item': 5},
         {'item': 6}
       ]);
     });

  it('can be created by zipping an array of datasets with primitive ' +
         'elements',
     async () => {
       const a = tfd.array([1, 2, 3]);
       const b = tfd.array([4, 5, 6]);
       const result = await tfd.zip([a, b]).toArrayForTest();
       expect(result).toEqual([[1, 4], [2, 5], [3, 6]]);
     });

  it('can be created by zipping an array of datasets with object elements',
     async () => {
       const a = tfd.array([{a: 1}, {a: 2}, {a: 3}]);
       const b = tfd.array([{b: 4}, {b: 5}, {b: 6}]);
       const result = await tfd.zip([a, b]).toArrayForTest();
       expect(result).toEqual(
           [[{a: 1}, {b: 4}], [{a: 2}, {b: 5}], [{a: 3}, {b: 6}]]);
     });

  it('can be created by zipping a dict of datasets', async () => {
    const a = tfd.array([{a: 1}, {a: 2}, {a: 3}]);
    const b = tfd.array([{b: 4}, {b: 5}, {b: 6}]);
    const result = await tfd.zip({c: a, d: b}).toArrayForTest();
    expect(result).toEqual([
      {c: {a: 1}, d: {b: 4}}, {c: {a: 2}, d: {b: 5}}, {c: {a: 3}, d: {b: 6}}
    ]);
  });

  it('can be created by zipping a nested structure of datasets', async () => {
    const a = tfd.array([1, 2, 3]);
    const b = tfd.array([4, 5, 6]);
    const c = tfd.array([7, 8, 9]);
    const d = tfd.array([10, 11, 12]);
    const result = await tfd.zip({a, bcd: [b, {c, d}]}).toArrayForTest();

    expect(result).toEqual([
      {a: 1, bcd: [4, {c: 7, d: 10}]},
      {a: 2, bcd: [5, {c: 8, d: 11}]},
      {a: 3, bcd: [6, {c: 9, d: 12}]},
    ]);
  });

  it('can be created by zipping datasets of different sizes', async () => {
    const a = tfd.array([1, 2]);
    const b = tfd.array([3, 4, 5, 6]);
    const result = await tfd.zip([a, b]).toArrayForTest();
    expect(result).toEqual([[1, 3], [2, 4]]);
  });

  it('zipping a native string throws an error', async done => {
    try {
      // tslint:disable-next-line:no-any no-construct
      await tfd.zip('test' as any);
      done.fail();
    } catch (e) {
      expect(e.message).toEqual(
          'The argument to zip() must be an object or array.');
      done();
    }
  });

  it('zipping a string object throws a meaningful error', async done => {
    try {
      // tslint:disable-next-line:no-any no-construct
      await tfd.zip(new String('test') as any).iterator();
      done.fail();
    } catch (e) {
      // This error is not specific to the error case arising from
      //   typeof(new String('test')) === 'object'
      // Instead this error is thrown because the leaves of the structure
      // are the letters t, e, s, and t, as well a number for the length. I
      // think it's a fine error message for this situation anyway.
      expect(e.message).toEqual(
          'Leaves of the structure passed to zip() must be Datasets, ' +
          'not primitives.');
      done();
    }
  });

  it('zipping a structure with repeated elements works', async () => {
    const a = tfd.array([1, 2, 3]);
    const b = tfd.array([4, 5, 6]);
    const c = tfd.array([7, 8, 9]);
    const d = tfd.array([10, 11, 12]);
    const result =
        await tfd.zip({a, abacd: [a, b, {a, c, d}]}).toArrayForTest();

    expect(result).toEqual([
      {a: 1, abacd: [1, 4, {a: 1, c: 7, d: 10}]},
      {a: 2, abacd: [2, 5, {a: 2, c: 8, d: 11}]},
      {a: 3, abacd: [3, 6, {a: 3, c: 9, d: 12}]},
    ]);
  });

  it('zipping a structure with cycles throws an error', async done => {
    try {
      // tslint:disable-next-line:no-any
      const a = tfd.array([1, 2, 3]);
      const b = tfd.array([4, 5, 6]);
      const c: DatasetContainer = [tfd.array([7, 8, 9])];
      const abc: DatasetContainer = [a, b, c];
      c.push(abc);
      await tfd.zip({a, abc}).iterator();
      done.fail();
    } catch (e) {
      expect(e.message).toEqual('Circular references are not supported.');
      done();
    }
  });

  it('zip propagates errors thrown when iterating constituent datasets',
     async done => {
       try {
         const makeIterator = () => {
           let count = 0;
           return {
             next: () => {
               if (count > 2) {
                 throw new Error('propagate me!');
               }
               return {value: count++, done: false};
             }
           };
         };
         const a = tfd.generator(makeIterator);
         const b = tfd.array([3, 4, 5, 6]);
         // Using toArray() rather than toArrayForTest().  The prefetch in
         // the latter, in combination with expecting an exception, causes
         // unrelated tests to fail (See
         // https://github.com/tensorflow/tfjs/issues/1330.
         await (await tfd.zip([a, b]).iterator()).toArray();
         done.fail();
       } catch (e) {
         expect(e.message).toMatch(/propagate me!/);
         done();
       }
     });

  it('can be repeated a fixed number of times', async () => {
    const a = tfd.array([{'item': 1}, {'item': 2}, {'item': 3}]);
    const result = await a.repeat(4).toArrayForTest();
    expect(result).toEqual([
      {'item': 1},
      {'item': 2},
      {'item': 3},
      {'item': 1},
      {'item': 2},
      {'item': 3},
      {'item': 1},
      {'item': 2},
      {'item': 3},
      {'item': 1},
      {'item': 2},
      {'item': 3},
    ]);
  });

  it('can be repeated indefinitely', async () => {
    const a = tfd.array([{'item': 1}, {'item': 2}, {'item': 3}]);
    await a.repeat().take(234).toArrayForTest();
  });

  it('can be repeated with state in a closure', async () => {
    // This tests a tricky bug having to do with 'this' being set properly.
    // See
    // https://github.com/Microsoft/TypeScript/wiki/%27this%27-in-TypeScript

    class CustomDataset extends tfd.Dataset<{}> {
      state = {val: 1};
      async iterator() {
        const result = iteratorFromItems([
          {'item': this.state.val++}, {'item': this.state.val++},
          {'item': this.state.val++}
        ]);
        return result;
      }
    }
    const a = new CustomDataset();
    await a.repeat().take(1234).toArrayForTest();
  });

  it('can collect all items into memory', async () => {
    const ds = new TestDataset();
    const items = await ds.toArrayForTest();
    expect(items.length).toEqual(100);
    // The test dataset has 100 elements, each containing 2 Tensors.
    expect(tf.memory().numTensors).toEqual(200);
  });

  it('batches entries into column-oriented batches', async () => {
    const ds = new TestDataset();
    const bds = ds.batch(8);
    const batchIterator = await bds.iterator();
    const result = await batchIterator.toArrayForTest();

    expect(result.length).toEqual(13);
    result.slice(0, 12).forEach(batch => {
      const b = batch as TensorContainerObject;
      expect((b['number'] as tf.Tensor).shape).toEqual([8]);
      expect((b['numberArray'] as tf.Tensor).shape).toEqual([8, 3]);
      expect((b['Tensor'] as tf.Tensor).shape).toEqual([8, 3]);
      expect((b['Tensor2'] as tf.Tensor).shape).toEqual([8, 2, 2]);
      expect((b['string'] as tf.Tensor).shape).toEqual([8]);
    });
    tf.dispose(result);
    expect(tf.memory().numTensors).toBe(0);
  });

  it('batches entries without leaking Tensors', async () => {
    // The prior test confirms this too, but this formulation is more
    // specific.

    // First show that an unbatched test iterator creates 2 Tensors per
    // element.
    const ds = new TestDataset();
    const iter = await ds.iterator();
    const element = await iter.next();
    expect(tf.memory().numTensors).toBe(2);
    tf.dispose(element.value);
    expect(tf.memory().numTensors).toBe(0);

    // Now obtain batches, which should contain 4 Tensors each.
    const bDs = new TestDataset().batch(8);
    const bIter = await bDs.iterator();
    const bElement = await bIter.next();
    // The batch element contains four Tensors, and the 8*2 Tensors in the
    // original unbatched elements have already been disposed.
    expect(tf.memory().numTensors).toBe(5);
    tf.dispose(bElement.value);
    expect(tf.memory().numTensors).toBe(0);
  });

  it('batches complex nested objects into column-oriented batches',
     async () => {
       // Our "complexified" examples map the simple examples into a deep
       // nested structure. This test shows that batching complexified
       // examples produces the same result as complexifying batches of
       // simple examples.

       const complexThenBatch =
           new TestDataset().map(complexifyExampleAsDict).batch(8);
       const batchThenComplex =
           new TestDataset().batch(8).map(complexifyExampleAsDict);

       const compareDataset = tfd.zip({complexThenBatch, batchThenComplex});

       const result = await (await compareDataset.iterator()).toArrayForTest();

       expect(result.length).toEqual(13);
       // tslint:disable-next-line:no-any
       result.slice(0, 12).forEach((compare: any) => {
         // TODO(soergel): could use deepMap to implement deep compare.
         // For now, just spot-check a few deep entries.
         expect(compare.complexThenBatch.a.v.shape)
             .toEqual(compare.batchThenComplex.a.v.shape);
         expect(compare.complexThenBatch.a.v.dataSync())
             .toEqual(compare.batchThenComplex.a.v.dataSync());

         expect(compare.complexThenBatch.a.q.ab.x.shape)
             .toEqual(compare.batchThenComplex.a.q.ab.x.shape);
         expect(compare.complexThenBatch.a.q.ab.x.dataSync())
             .toEqual(compare.batchThenComplex.a.q.ab.x.dataSync());

         expect(compare.complexThenBatch.b.bb.bbb.shape)
             .toEqual(compare.batchThenComplex.b.bb.bbb.shape);
         expect(compare.complexThenBatch.b.bb.bbb.dataSync())
             .toEqual(compare.batchThenComplex.b.bb.bbb.dataSync());

         expect(compare.complexThenBatch.c.ca.x.shape)
             .toEqual(compare.batchThenComplex.c.ca.x.shape);
         expect(compare.complexThenBatch.c.ca.x.dataSync())
             .toEqual(compare.batchThenComplex.c.ca.x.dataSync());
       });
       tf.dispose(result);
       expect(tf.memory().numTensors).toBe(0);
     });

  it('batches nested numeric arrays into a single Tensor', async () => {
    const dataset =
        new TestDataset()
            .map((e) => {
              const a = e.number as number;
              const b = a * 2;
              const c = b + 5;
              return {
                foo: [[[a, b], [c, a], [b, c]], [[b, c], [a, b], [c, a]]]
              };
            })
            .batch(8);

    const result = await (await dataset.iterator()).toArrayForTest();
    expect(result.length).toEqual(13);

    // tslint:disable-next-line:no-any
    result.slice(0, 12).forEach((e: any) => {
      expect(e.foo instanceof tf.Tensor).toBeTruthy();
      expect(e.foo.shape).toEqual([8, 2, 3, 2]);
    });
    // tslint:disable-next-line:no-any
    expect((result[12] as any).foo.shape).toEqual([4, 2, 3, 2]);

    tf.dispose(result);
    expect(tf.memory().numTensors).toBe(0);
  });

  // TODO(soergel, smilkov): Reinstate this once tfjs-core enforces
  // TensorLike.
  /*
  it('throws an error when given an array containing a dict', async done
  => {
    const dataset = array([[1, {a: 2, b: 3}], [4, {a: 5, b: 6}]]).batch(2);
    try {
      await (await dataset.iterator()).collect();
      done.fail();
    } catch (e) {
      expect(e.message).toEqual('TODO');
      done();
    }
    expect(tf.memory().numTensors).toBe(0);
  });
  */

  it('batch creates a small last batch', async () => {
    const ds = new TestDataset();
    const bds = ds.batch(8);
    const batchIterator = await bds.iterator();
    const result = await batchIterator.toArrayForTest();
    const lastBatch = result[result.length - 1] as TensorContainerObject;
    expect((lastBatch['number'] as tf.Tensor).shape).toEqual([4]);
    expect((lastBatch['numberArray'] as tf.Tensor).shape).toEqual([4, 3]);
    expect((lastBatch['Tensor'] as tf.Tensor).shape).toEqual([4, 3]);
    expect((lastBatch['Tensor2'] as tf.Tensor).shape).toEqual([4, 2, 2]);
    expect((lastBatch['string'] as tf.Tensor).shape).toEqual([4]);

    const expectedNumberLastBatch = tf.tensor1d([96, 97, 98, 99]);
    tf.test_util.expectArraysClose(
        await (lastBatch['number'] as tf.Tensor).array(),
        await expectedNumberLastBatch.array());

    const expectedNumberArrayLastBatch = tf.tensor2d(
        [
          [96, 96 ** 2, 96 ** 3], [97, 97 ** 2, 97 ** 3],
          [98, 98 ** 2, 98 ** 3], [99, 99 ** 2, 99 ** 3]
        ],
        [4, 3]);
    tf.test_util.expectArraysClose(
        await (lastBatch['numberArray'] as tf.Tensor).array(),
        await expectedNumberArrayLastBatch.array());

    const expectedTensorLastBatch = tf.tensor2d(
        [
          [96, 96 ** 2, 96 ** 3], [97, 97 ** 2, 97 ** 3],
          [98, 98 ** 2, 98 ** 3], [99, 99 ** 2, 99 ** 3]
        ],
        [4, 3]);
    tf.test_util.expectArraysClose(
        await (lastBatch['Tensor'] as tf.Tensor).array(),
        await expectedTensorLastBatch.array());

    const expectedTensor2LastBatch = tf.tensor3d(
        [
          [[96, 96 ** 2], [96 ** 3, 96 ** 4]],
          [[97, 97 ** 2], [97 ** 3, 97 ** 4]],
          [[98, 98 ** 2], [98 ** 3, 98 ** 4]],
          [[99, 99 ** 2], [99 ** 3, 99 ** 4]],
        ],
        [4, 2, 2]);
    tf.test_util.expectArraysClose(
        await (lastBatch['Tensor2'] as tf.Tensor).array(),
        await expectedTensor2LastBatch.array());

    const expectedStringLastBatch =
        tf.tensor1d(['Item 96', 'Item 97', 'Item 98', 'Item 99']);
    tf.test_util.expectArraysEqual(
        await (lastBatch['string'] as tf.Tensor).array(),
        await expectedStringLastBatch.array());

    tf.dispose(result);
    tf.dispose(expectedNumberLastBatch);
    tf.dispose(expectedNumberArrayLastBatch);
    tf.dispose(expectedTensorLastBatch);
    tf.dispose(expectedTensor2LastBatch);
    tf.dispose(expectedStringLastBatch);

    expect(tf.memory().numTensors).toBe(0);
  });

  it('skip does not leak Tensors', async done => {
    try {
      const ds = new TestDataset();
      expect(tf.memory().numTensors).toEqual(0);
      const result = await ds.skip(15).toArrayForTest();
      // The test dataset had 100 elements; we skipped 15; 85 remain.
      expect(result.length).toEqual(85);
      // Each element of the test dataset contains 2 Tensors;
      // 85 elements remain, so 2 * 85 = 170 Tensors remain.
      expect(tf.memory().numTensors).toEqual(170);
      done();
    } catch (e) {
      done.fail(e);
    }
  });

  it('filter does not leak Tensors', async () => {
    const ds = new TestDataset();
    expect(tf.memory().numTensors).toEqual(0);
    await ds.filter(x => ((x['number'] as number) % 2 === 0)).toArrayForTest();
    // Each element of the test dataset contains 2 Tensors.
    // There were 100 elements, but we filtered out half of them.
    // Thus 50 * 2 = 100 Tensors remain.
    expect(tf.memory().numTensors).toEqual(100);
  });

  it('shuffle does not leak Tensors', async () => {
    const ds = new TestDataset();
    expect(tf.memory().numTensors).toEqual(0);
    await ds.shuffle(1000).toArrayForTest();
    // The shuffle operation emitted all of the tensors.
    expect(tf.memory().numTensors).toEqual(200);
  });

  it('shuffle throws an error when bufferSize is not specified and ' +
         'dataset.size is unknown.',
     async () => {
       const ds = new TestDataset();

       expect(() => ds.shuffle(undefined))
           .toThrowError(
               '`Dataset.shuffle()` requires bufferSize to be specified.');
     });

  it('shuffle throws an error when bufferSize is not specified and ' +
         'dataset.size is known.',
     async () => {
       const ds = new TestDataset(true);

       expect(() => ds.shuffle(undefined))
           .toThrowError(
               '`Dataset.shuffle()` requires bufferSize to be ' +
               'specified.  If your data fits in main memory (for ' +
               'regular JS objects), and/or GPU memory (for ' +
               '`tf.Tensor`s), consider setting bufferSize to the ' +
               'dataset size (200 elements)');
     });

  it('prefetch throws an error when bufferSize is not specified.', async () => {
    const ds = new TestDataset();

    expect(() => ds.prefetch(undefined))
        .toThrowError(
            '`Dataset.prefetch()` requires bufferSize to be specified.');
  });

  it('prefetch does not leak Tensors', async () => {
    const ds = new TestDataset();
    expect(tf.memory().numTensors).toEqual(0);
    await ds.prefetch(1000).toArray();
    // The prefetch operation emitted all of the tensors.
    expect(tf.memory().numTensors).toEqual(200);
  });

  it('map does not leak Tensors when none are returned', async () => {
    const ds = new TestDataset();
    expect(tf.memory().numTensors).toEqual(0);
    await ds.map(x => ({'constant': 1})).toArrayForTest();
    // The map operation consumed all of the tensors and emitted none.
    expect(tf.memory().numTensors).toEqual(0);
  });

  it('map does not lose or leak Tensors when some inputs are passed ' +
         'through',
     async () => {
       const ds = new TestDataset();
       expect(tf.memory().numTensors).toEqual(0);
       await ds.map(x => ({'Tensor2': x['Tensor2']})).toArrayForTest();
       // Each element of the test dataset contains 2 Tensors.
       // Our map operation retained one of the Tensors and discarded the
       // other. Thus the mapped data contains 100 elements with 1 Tensor
       // each.
       expect(tf.memory().numTensors).toEqual(100);
     });

  it('map does not leak Tensors when inputs are replaced', async () => {
    const ds = new TestDataset();
    expect(tf.memory().numTensors).toEqual(0);
    await ds.map(x => ({'a': tf.tensor1d([1, 2, 3])})).toArrayForTest();
    // Each element of the test dataset contains 2 Tensors.
    // Our map operation discarded both Tensors and created one new one.
    // Thus the mapped data contains 100 elements with 1 Tensor each.
    expect(tf.memory().numTensors).toEqual(100);
  });

  it('forEach does not leak Tensors', async () => {
    const ds = new TestDataset();
    let count = 0;
    await ds.forEachAsync(element => {
      count++;
      return {};
    });
    // forEach traversed the entire dataset of 100 elements.
    expect(count).toEqual(100);
    // forEach consumed all of the input Tensors.
    expect(tf.memory().numTensors).toEqual(0);
  });

  it('clone tensors when returning iterator of a dataset generated from ' +
         'existing tensors',
     async () => {
       expect(tf.memory().numTensors).toEqual(0);
       const a = tf.ones([2, 1]);
       const b = tf.ones([2, 1]);
       expect(tf.memory().numTensors).toEqual(2);
       const ds = tfd.array([a, b]);
       // Pre-existing tensors are not cloned during dataset creation.
       expect(tf.memory().numTensors).toEqual(2);

       let count = 0;
       // ds.forEachAsync() automatically disposes incoming Tensors after
       // processing them.
       await ds.forEachAsync(elem => {
         count++;
         expect(elem.isDisposed).toBeFalsy();
       });
       expect(count).toEqual(2);
       // Cloned tensors are disposed after traverse, while original tensors
       // stay.
       expect(tf.memory().numTensors).toEqual(2);

       await ds.forEachAsync(elem => {
         count++;
         expect(elem.isDisposed).toBeFalsy();
       });
       expect(count).toEqual(4);
       expect(tf.memory().numTensors).toEqual(2);

       await ds.forEachAsync(elem => {
         count++;
         expect(elem.isDisposed).toBeFalsy();
       });
       expect(count).toEqual(6);
       expect(tf.memory().numTensors).toEqual(2);

       expect(a.isDisposed).toBeFalsy();
       expect(b.isDisposed).toBeFalsy();
     });

  it('clone tensors in nested structures when returning iterator of a ' +
         'dataset generated from existing tensors',
     async () => {
       expect(tf.memory().numTensors).toEqual(0);
       const a = tf.ones([2, 1]);
       const b = tf.ones([2, 1]);
       const c = tf.ones([2, 1]);
       const d = tf.ones([2, 1]);
       expect(tf.memory().numTensors).toEqual(4);
       const ds = tfd.array([{foo: a, bar: b}, {foo: c, bar: d}]);
       // Pre-existing tensors are not cloned during dataset creation.
       expect(tf.memory().numTensors).toEqual(4);

       let count = 0;
       // ds.forEachAsync() automatically disposes incoming Tensors after
       // processing them.
       await ds.forEachAsync(elem => {
         count++;
         expect(elem.foo.isDisposed).toBeFalsy();
         expect(elem.bar.isDisposed).toBeFalsy();
       });
       expect(count).toEqual(2);
       // Cloned tensors are disposed after traverse, while original tensors
       // stay.
       expect(tf.memory().numTensors).toEqual(4);

       await ds.forEachAsync(elem => {
         count++;
         expect(elem.foo.isDisposed).toBeFalsy();
         expect(elem.bar.isDisposed).toBeFalsy();
       });
       expect(count).toEqual(4);
       expect(tf.memory().numTensors).toEqual(4);

       await ds.forEachAsync(elem => {
         count++;
         expect(elem.foo.isDisposed).toBeFalsy();
         expect(elem.bar.isDisposed).toBeFalsy();
       });
       expect(count).toEqual(6);
       expect(tf.memory().numTensors).toEqual(4);

       expect(a.isDisposed).toBeFalsy();
       expect(b.isDisposed).toBeFalsy();
     });

  it('traverse dataset from tensors without leaking Tensors', async () => {
    expect(tf.memory().numTensors).toEqual(0);
    const a = tf.ones([2, 1]);
    const b = tf.ones([2, 1]);
    const c = tf.ones([2, 1]);
    const d = tf.ones([2, 1]);
    expect(tf.memory().numTensors).toEqual(4);
    const ds = tfd.array([a, b, c, d]).take(2);
    // Pre-existing tensors are not cloned during dataset creation.
    expect(tf.memory().numTensors).toEqual(4);

    let count = 0;
    // ds.forEachAsync() automatically disposes incoming Tensors after
    // processing them.
    await ds.forEachAsync(elem => {
      count++;
      expect(elem.isDisposed).toBeFalsy();
    });
    expect(count).toEqual(2);
    // Cloned tensors are disposed after traverse, while original tensors
    // stay.
    expect(tf.memory().numTensors).toEqual(4);

    await ds.forEachAsync(elem => {
      count++;
      expect(elem.isDisposed).toBeFalsy();
    });
    expect(count).toEqual(4);
    expect(tf.memory().numTensors).toEqual(4);

    await ds.forEachAsync(elem => {
      count++;
      expect(elem.isDisposed).toBeFalsy();
    });
    expect(count).toEqual(6);
    expect(tf.memory().numTensors).toEqual(4);

    expect(a.isDisposed).toBeFalsy();
    expect(b.isDisposed).toBeFalsy();
  });

  it('can get correct size of dataset from objects array', async () => {
    const ds = tfd.array([{'item': 1}, {'item': 2}, {'item': 3}]);
    expect(ds.size).toEqual(3);
  });

  it('can get correct size of dataset from number array', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]);
    expect(ds.size).toEqual(5);
  });

  it('can get size 0 from empty dataset', async () => {
    const ds = tfd.array([]);
    expect(ds.size).toEqual(0);
  });

  it('size is null if dataset may exhausted randomly', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator);
    expect(ds.size).toBeNull();
  });

  it('repeat dataset has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat(3);
    expect(ds.size).toEqual(15);
  });

  it('repeat dataset forever has infinite size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat();
    expect(ds.size).toEqual(Infinity);
  });

  it('repeat unknown size dataset has null size', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator).repeat(3);
    expect(ds.size).toBeNull();
  });

  it('take dataset has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).take(3);
    expect(ds.size).toEqual(3);
  });

  it('take dataset without enough elements has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).take(10);
    expect(ds.size).toEqual(5);
  });

  it('take dataset with unknown size has null size', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator).take(3);
    expect(ds.size).toBeNull();
  });

  it('take dataset with infinite elements has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat().take(10);
    expect(ds.size).toEqual(10);
  });

  it('skip dataset has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).skip(2);
    expect(ds.size).toEqual(3);
  });

  it('skip dataset without enough elements has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).skip(10);
    expect(ds.size).toEqual(0);
  });

  it('skip dataset with unknown size has null size', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator).skip(3);
    expect(ds.size).toBeNull();
  });

  it('skip dataset with infinite elements has infinity size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat().skip(10);
    expect(ds.size).toEqual(Infinity);
  });

  it('batch dataset with small last batch has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5, 6, 7]).batch(2, true);
    expect(ds.size).toEqual(4);
  });

  it('batch dataset without small last batch has correct size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5, 6, 7]).batch(2, false);
    expect(ds.size).toEqual(3);
  });

  it('batch dataset with unknown size has null size', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator).batch(2);
    expect(ds.size).toBeNull();
  });

  it('batch dataset with infinite elements has infinity size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat().batch(2);
    expect(ds.size).toEqual(Infinity);
  });

  it('map dataset preserves regular size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).map(e => e + 1);
    expect(ds.size).toEqual(5);
  });

  it('map dataset preserves infinity size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat().map(e => e + 1);
    expect(ds.size).toEqual(Infinity);
  });

  it('map dataset preserves null size', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator).map(e => e + 1);
    expect(ds.size).toBeNull();
  });

  it('filter dataset preserves infinity size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).repeat().filter(e => e % 2 === 0);
    expect(ds.size).toEqual(Infinity);
  });

  it('filter dataset with regular size has null size', async () => {
    const ds = tfd.array([1, 2, 3, 4, 5]).filter(e => e % 2 === 0);
    expect(ds.size).toBeNull();
  });

  it('filter dataset with null size has null size', async () => {
    const makeIterator = () => {
      let i = -1;
      return {
        next: () =>
            ++i < 7 ? {value: i, done: false} : {value: null, done: true}
      };
    };
    const ds = tfd.generator(makeIterator).filter(e => e % 2 === 0);
    expect(ds.size).toBeNull();
  });

  it('zipping an array of datasets with primitive elements has correct ' +
         'size',
     async () => {
       const a = tfd.array([1, 2, 3]);
       const b = tfd.array([4, 5, 6, 7, 8]);
       const result = await tfd.zip([a, b]);
       expect(result.size).toEqual(3);
     });

  it('zipping an array of datasets with object elements has correct size',
     async () => {
       const a = tfd.array([{a: 1}, {a: 2}, {a: 3}]);
       const b = tfd.array([{b: 4}, {b: 5}, {b: 6}, {b: 7}, {b: 8}]);
       const result = await tfd.zip([a, b]);
       expect(result.size).toEqual(3);
     });

  it('zipping an object of datasets with primitive elements has correct ' +
         'size',
     async () => {
       const a = tfd.array([1, 2, 3]);
       const b = tfd.array([4, 5, 6, 7, 8]);
       const result = await tfd.zip({'a': a, 'b': b});
       expect(result.size).toEqual(3);
     });

  it('zipping an object of datasets with object elements has correct size',
     async () => {
       const a = tfd.array([{a: 1}, {a: 2}, {a: 3}]);
       const b = tfd.array([{b: 4}, {b: 5}, {b: 6}, {b: 7}, {b: 8}]);
       const result = await tfd.zip({'a': a, 'b': b});
       expect(result.size).toEqual(3);
     });

  it('converting dataset with infinite elements to array throws error',
     async done => {
       try {
         const ds = tfd.array([1, 2, 3, 4, 5]).repeat();
         expect(ds.size).toEqual(Infinity);
         await ds.toArrayForTest();
         done.fail();
       } catch (e) {
         expect(e.message).toEqual(
             'Can not convert infinite data stream to array.');
         done();
       }
     });

  it('ArrayBuffer type data is not converted', async () => {
    const a = new Float32Array([1, 2, 3]);
    const ds = tfd.array([a]);
    const result = await ds.toArrayForTest();
    const resultArray = result[0];
    expect(resultArray[0]).toBe(1);
    expect(resultArray[1]).toBe(2);
    expect(resultArray[2]).toBe(3);
    expect(ArrayBuffer.isView(resultArray)).toBe(true);
    expect(resultArray.constructor === Float32Array).toBe(true);
  });
});

describeAllEnvs('Dataset with DEBUG mode', () => {
  beforeAll(() => {
    tf.env().set('DEBUG', true);
  });

  it('throws an error when given an array of inconsistent shape',
     async done => {
       const dataset = array([[[1, 2], [3]], [[4, 5], [6]]]).batch(2);
       try {
         // Using toArray() rather than toArrayForTest().  The prefetch in
         // the latter, in combination with expecting an exception, causes
         // unrelated tests to fail (See
         // https://github.com/tensorflow/tfjs/issues/1330.
         await (await dataset.iterator()).toArray();
         done.fail();
       } catch (e) {
         expect(e.message).toEqual(
             'Element arr[0][1] should have 2 elements, ' +
             'but has 1 elements');
         done();
       }
       expect(tf.memory().numTensors).toBe(0);
     });
});