license import * as tf from '@tensorflow/tfjs-core' ; import { TensorLike } from '@tensorflow/tfjs-core' ; import * as seedrandom from 'seedrandom' ; import { iteratorFromConcatenated , iteratorFromFunction , iteratorFromItems , iteratorFromZipped , LazyIterator , ZipMismatchMode } from './iterators/lazy_iterator' ; import { DatasetContainer } from './types' ; import { canTensorify , deepMapAndAwaitAll , DeepMapResult , isIterable } from './util/deep_map' ; /**
 * Represents a potentially large list of independent data elements (typically
 * 'samples' or 'examples').
 *
 * A 'data example' may be a primitive, an array, a map from string keys to
 * values, or any nested structure of these.
 *
 * A `Dataset` represents an ordered collection of elements, together with a
 * chain of transformations to be performed on those elements. Each
 * transformation is a method of `Dataset` that returns another `Dataset`, so
 * these may be chained, e.g.
 * `const processedDataset = rawDataset.filter(...).map(...).batch(...)`.
 *
 * Data loading and transformation is done in a lazy, streaming fashion.  The
 * dataset may be iterated over multiple times; each iteration starts the data
 * loading anew and recapitulates the transformations.
 *
 * A `Dataset` is typically processed as a stream of unbatched examples --i.e.,
 * its transformations are applied one example at a time. Batching produces a
 * new `Dataset` where each element is a batch. Batching should usually come
 * last in a pipeline, because data transformations are easier to express on a
 * per-example basis than on a per-batch basis.
 *
 * The following code examples are calling `await dataset.forEachAsync(...)` to
 * iterate once over the entire dataset in order to print out the data.
 */ doc export abstract class Dataset < T extends tf . TensorContainer > { abstract async iterator ( ) ; readonly size : number = null ; param batchSize param smallLastBatch returns doc batch ( batchSize , smallLastBatch = true ) { const base = this ; tf . util . assert ( batchSize > 0 , ( ) => `batchSize needs to be positive, but it is
      ${ batchSize }` ) ; let size ; if ( this . size === Infinity || this . size == null ) { size = this . size ; } else if ( smallLastBatch ) { size = Math . ceil ( this . size / batchSize ) ; } else { size = Math . floor ( this . size / batchSize ) ; } return datasetFromIteratorFn ( async ( ) => { return ( await base . iterator ( ) ) . columnMajorBatch ( batchSize , smallLastBatch , deepBatchConcat ) ; } , size ) ; } param dataset returns doc concatenate ( dataset ) { const base = this ; let size ; if ( this . size === Infinity || dataset . size === Infinity ) { size = Infinity ; } else if ( this . size != null && dataset . size != null ) { size = this . size + dataset . size ; } else { size = null ; } return datasetFromIteratorFn ( async ( ) => ( await base . iterator ( ) ) . concatenate ( await dataset . iterator ( ) ) , size ) ; } param predicate returns doc filter ( predicate ) { const base = this ; let size ; if ( this . size === Infinity ) { size = Infinity ; } else { size = null ; } return datasetFromIteratorFn ( async ( ) => { return ( await base . iterator ( ) ) . filter ( x => tf . tidy ( ( ) => predicate ( x ) ) ) ; } , size ) ; } param f returns doc async forEachAsync ( f ) { return ( await this . iterator ( ) ) . forEachAsync ( f ) ; } deprecated async forEach ( f ) { tf . deprecationWarn ( 'dataset.forEach() is deprecated and will be removed. ' + 'Please use dataset.forEachAsync() instead' ) ; return this . forEachAsync ( f ) ; } param transform returns doc map < O extends tf . TensorContainer > ( transform ) { const base = this ; return datasetFromIteratorFn ( async ( ) => { return ( await base . iterator ( ) ) . map ( x => tf . tidy ( ( ) => transform ( x ) ) ) ; } , this . size ) ; } param transform returns doc mapAsync < O extends tf . TensorContainer > ( transform ) { const base = this ; return datasetFromIteratorFn ( async ( ) => { return ( await base . iterator ( ) ) . mapAsync ( transform ) ; } , this . size ) ; } param bufferSize returns doc prefetch ( bufferSize ) { if ( bufferSize == null ) { throw new RangeError ( '`Dataset.prefetch()` requires bufferSize to be specified.' ) ; } const base = this ; return datasetFromIteratorFn ( async ( ) => ( await base . iterator ( ) ) . prefetch ( bufferSize ) , this . size ) ; } param count returns doc repeat ( count ? ) { const base = this ; let size ; if ( this . size != null && count > 0 ) { size = this . size * count ; } else if ( count === 0 ) { size = 0 ; } else if ( this . size != null && ( count === undefined || count < 0 ) ) { size = Infinity ; } else { size = null ; } return datasetFromIteratorFn ( async ( ) => { const iteratorIterator = iteratorFromFunction ( async ( ) => ( { value : await base . iterator ( ) , done : false } ) ) ; return iteratorFromConcatenated ( iteratorIterator . take ( count ) ) ; } , size ) ; } param count returns doc skip ( count ) { const base = this ; let size ; if ( this . size != null && count >= 0 && this . size >= count ) { size = this . size - count ; } else if ( this . size != null && ( this . size < count || count === undefined || count < 0 ) ) { size = 0 ; } else { size = null ; } return datasetFromIteratorFn ( async ( ) => ( await base . iterator ( ) ) . skip ( count ) , size ) ; } static readonly MAX_BUFFER_SIZE = 10000 ; param bufferSize param seed param reshuffleEachIteration returns doc shuffle ( bufferSize , seed ? , reshuffleEachIteration = true ) { if ( bufferSize == null || bufferSize < 0 ) { if ( this . size == null ) { throw new RangeError ( '`Dataset.shuffle()` requires bufferSize to be specified.' ) ; } else { throw new RangeError ( '`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 (${ this . size } elements)` ) ; } } const base = this ; const random = seedrandom . alea ( seed || tf . util . now ( ) . toString ( ) ) ; return datasetFromIteratorFn ( async ( ) => { let seed2 = random . int32 ( ) ; if ( reshuffleEachIteration ) { seed2 += random . int32 ( ) ; } return ( await base . iterator ( ) ) . shuffle ( bufferSize , seed2 . toString ( ) ) ; } , this . size ) ; } param count returns doc take ( count ) { const base = this ; let size ; if ( this . size != null && this . size > count ) { size = count ; } else if ( this . size != null && this . size <= count ) { size = this . size ; } else { size = null ; } return datasetFromIteratorFn ( async ( ) => ( await base . iterator ( ) ) . take ( count ) , size ) ; } returns doc async toArray ( ) { if ( this . size === Infinity ) { throw new Error ( 'Can not convert infinite data stream to array.' ) ; } return ( await this . iterator ( ) ) . toArray ( ) ; } returns async toArrayForTest ( ) { if ( this . size === Infinity ) { throw new Error ( 'Can not convert infinite data stream to array.' ) ; } return ( await this . iterator ( ) ) . toArrayForTest ( ) ; } } /**
 * Create a `Dataset` defined by a provided iterator() function.
 *
 * ```js
 * let i = -1;
 * const func = () =>
 *    ++i < 5 ? {value: i, done: false} : {value: null, done: true};
 * const iter = tf.data.iteratorFromFunction(func);
 * const ds = tf.data.datasetFromIteratorFn(iter);
 * await ds.forEachAsync(e => console.log(e));
 * ```
 */ export function datasetFromIteratorFn < T extends tf . TensorContainer > ( iteratorFn , size = null ) { return new class extends Dataset < T > { size = size ; async iterator ( ) { return iteratorFn ( ) ; } } ( ) ; } param items doc export function array < T extends tf . TensorContainer > ( items ) { return datasetFromIteratorFn ( async ( ) => iteratorFromItems ( items ) , items . length ) ; } /**
 * Create a `Dataset` by zipping together an array, dict, or nested
 * structure of `Dataset`s (and perhaps additional constants).
 * The underlying datasets must provide elements in a consistent order such that
 * they correspond.
 *
 * The number of elements in the resulting dataset is the same as the size of
 * the smallest dataset in datasets.
 *
 * The nested structure of the `datasets` argument determines the
 * structure of elements in the resulting iterator.
 *
 * Note this means that, given an array of two datasets that produce dict
 * elements, the result is a dataset that produces elements that are arrays
 * of two dicts:
 *
 * Zip an array of datasets:
 * ```js
 * console.log('Zip two datasets of objects:');
 * const ds1 = tf.data.array([{a: 1}, {a: 2}, {a: 3}]);
 * const ds2 = tf.data.array([{b: 4}, {b: 5}, {b: 6}]);
 * const ds3 = tf.data.zip([ds1, ds2]);
 * await ds3.forEachAsync(e => console.log(JSON.stringify(e)));
 *
 * // If the goal is to merge the dicts in order to produce elements like
 * // {a: ..., b: ...}, this requires a second step such as:
 * console.log('Merge the objects:');
 * const ds4 = ds3.map(x => {return {a: x[0].a, b: x[1].b}});
 * await ds4.forEachAsync(e => console.log(e));
 * ```
 *
 * Zip a dict of datasets:
 * ```js
 * const a = tf.data.array([{a: 1}, {a: 2}, {a: 3}]);
 * const b = tf.data.array([{b: 4}, {b: 5}, {b: 6}]);
 * const c = tf.data.zip({c: a, d: b});
 * await c.forEachAsync(e => console.log(JSON.stringify(e)));
 * ```
 */ doc export function zip < O extends tf . TensorContainer > ( datasets ) { if ( ! isIterable ( datasets ) ) { throw new Error ( 'The argument to zip() must be an object or array.' ) ; } let size ; if ( Array . isArray ( datasets ) ) { for ( let i = 0 ; i < datasets . length ; i ++ ) { size = size == null ? ( datasets [ i ] as Dataset < O > ) . size : Math . min ( size , ( datasets [ i ] as Dataset < O > ) . size ) ; } } else if ( datasets instanceof Object ) { for ( const ds in datasets ) { size = size == null ? ( datasets [ ds ] as Dataset < O > ) . size : Math . min ( size , ( datasets [ ds ] as Dataset < O > ) . size ) ; } } return datasetFromIteratorFn < O > ( async ( ) => { const streams = await deepMapAndAwaitAll ( datasets , d => { if ( d instanceof Dataset ) { return { value : d . iterator ( ) , recurse : false } ; } else if ( isIterable ( d ) ) { return { value : null , recurse : true } ; } else { throw new Error ( 'Leaves of the structure passed to zip() must be Datasets, ' + 'not primitives.' ) ; } } ) ; return iteratorFromZipped < O > ( streams , ZipMismatchMode . SHORTEST ) ; } , size ) ; } /**
 * A zip function for use with deepZip, passed via the columnMajorBatch call.
 *
 * Accepts an array of identically-structured nested elements and either batches
 * them (if they are primitives, numeric arrays, or Tensors) or requests
 * recursion (if not).
 */ function deepBatchConcat ( rows ) { if ( rows === null ) { return null ; } const exampleRow = rows [ 0 ] ; if ( canTensorify ( exampleRow ) ) { const value = batchConcat ( rows ) ; return { value , recurse : false } ; } return { value : null , recurse : true } ; } /**
 * Assembles a list of same-shaped numbers, number arrays, or Tensors
 * into a single new Tensor where axis 0 is the batch dimension.
 */ function batchConcat < T extends ( TensorLike | tf . Tensor ) > ( arrays ) { if ( arrays . length === 0 ) { throw new Error ( 'Can\'t make a batch of zero elements.' ) ; } if ( arrays [ 0 ] instanceof tf . Tensor ) { return tf . stack ( arrays as tf . Tensor [ ] ) ; } else { return tf . tensor ( arrays as TensorLike ) ; } }