license import * as tf from '@tensorflow/tfjs-core' ; import * as seedrandom from 'seedrandom' ; import { Container } from '../types' ; import { deepClone } from '../util/deep_clone' ; import { deepMapAndAwaitAll , DeepMapAsyncResult , DeepMapResult , deepZip , zipToList } from '../util/deep_map' ; import { GrowingRingBuffer } from '../util/growing_ring_buffer' ; import { RingBuffer } from '../util/ring_buffer' ; /**
 * A nested structure of LazyIterators, used as the input to zip().
 */ export type IteratorContainer = Container < LazyIterator < tf . TensorContainer > > ; /**
 * Create a `LazyIterator` from an array of items.
 */ export function iteratorFromItems < T > ( items ) { return new ArrayIterator ( items ) ; } /**
 * Create a `LazyIterator` of incrementing integers.
 */ export function iteratorFromIncrementing ( start ) { let i = start ; return iteratorFromFunction ( ( ) => ( { value : i ++ , done : false } ) ) ; } param func export function iteratorFromFunction < T > ( func ) { return new FunctionCallIterator ( func ) ; } param baseIterators param baseErrorHandler export function iteratorFromConcatenated < T > ( baseIterators , baseErrorHandler ? ) { return new ChainedIterator ( baseIterators , baseErrorHandler ) ; } param iteratorFunc param count param baseErrorHandler export function iteratorFromConcatenatedFunction < T > ( iteratorFunc , count , baseErrorHandler ? ) { return iteratorFromConcatenated ( iteratorFromFunction ( iteratorFunc ) . take ( count ) , baseErrorHandler ) ; } param iterators param mismatchMode export function iteratorFromZipped < O extends tf . TensorContainer > ( iterators , mismatchMode = ZipMismatchMode . FAIL ) { return new ZipIterator < O > ( iterators , mismatchMode ) ; } /**
 * An asynchronous iterator, providing lazy access to a potentially
 * unbounded stream of elements.
 *
 * Iterator can be obtained from a dataset:
 * `const iter = await dataset.iterator();`
 */ export abstract class LazyIterator < T > { abstract summary ( ) ; /**
     * Returns a `Promise` for the next element in the stream.
     *
     * When an item can be provided successfully, the return value is
     * `{value:T, done:false}`.
     *
     * Calling next() on a closed stream returns `{value:null, done:true}`.
     */ abstract async next ( ) ; returns async toArray ( ) { const result = [ ] ; let x = await this . next ( ) ; while ( ! x . done ) { result . push ( x . value ) ; x = await this . next ( ) ; } return result ; } returns async toArrayForTest ( ) { const stream = this . prefetch ( 100 ) ; const result = [ ] ; let x = await stream . next ( ) ; while ( ! x . done ) { result . push ( x . value ) ; x = await stream . next ( ) ; } return result ; } /**
     * Draw items from the stream until it is exhausted.
     *
     * This can be useful when the stream has side effects but no output.  In
     * that case, calling this function guarantees that the stream will be
     * fully processed.
     */ async resolveFully ( ) { let x = await this . next ( ) ; while ( ! x . done ) { x = await this . next ( ) ; } } /**
     * Draw items from the stream until it is exhausted, or a predicate fails.
     *
     * This can be useful when the stream has side effects but no output.  In
     * that case, calling this function guarantees that the stream will be
     * fully processed.
     */ async resolveWhile ( predicate ) { let x = await this . next ( ) ; let shouldContinue = predicate ( x . value ) ; while ( ( ! x . done ) && shouldContinue ) { x = await this . next ( ) ; shouldContinue = predicate ( x . value ) ; } } param handler returns handleErrors ( handler ) { return new ErrorHandlingLazyIterator ( this , handler ) ; } param predicate returns filter ( predicate ) { return new FilterIterator ( this , predicate ) ; } param transform returns map < O > ( transform ) { return new MapIterator ( this , transform ) ; } param transform returns mapAsync < O > ( transform ) { return new AsyncMapIterator ( this , transform ) ; } param transform returns serialMapAsync < O > ( transform ) { return new AsyncMapIterator ( this , transform ) . serial ( ) ; } param transform returns flatmap < O > ( transform ) { return new FlatmapIterator ( this , transform ) ; } param f async forEachAsync ( f ) { return this . map ( f ) . resolveFully ( ) ; } param f async serialForEach ( f ) { return this . serialMapAsync ( f ) . resolveWhile ( x => ( x === true ) ) ; } param batchSize param smallLastBatch returns rowMajorBatch ( batchSize , smallLastBatch = true ) { return new RowMajorBatchIterator ( this , batchSize , smallLastBatch ) ; } param batchSize param smallLastBatch param zipFn returns columnMajorBatch ( batchSize , smallLastBatch = true , zipFn = zipToList ) { const rowBatches = this . rowMajorBatch ( batchSize , smallLastBatch ) ; return rowBatches . map ( x => deepZip ( x , zipFn ) ) ; } param iterator param baseErrorHandler returns concatenate ( iterator , baseErrorHandler ? ) { return new ChainedIterator ( iteratorFromItems ( [ this , iterator ] ) , baseErrorHandler ) ; } param count take ( count ) { if ( count < 0 || count == null ) { return this ; } return new TakeIterator ( this , count ) ; } param count skip ( count ) { if ( count < 0 || count == null ) { return this ; } return new SkipIterator ( this , count ) ; } param bufferSize prefetch ( bufferSize ) { return new PrefetchIterator ( this , bufferSize ) ; } param bufferSize param seed shuffle ( windowSize , seed ? ) { return new ShuffleIterator ( this , windowSize , seed ) ; } /**
     * Force an iterator to execute serially: each next() call will await the
     * prior one, so that they cannot execute concurrently.
     */ serial ( ) { return new SerialIterator ( this ) ; } } class ArrayIterator < T > extends LazyIterator < T > { private trav = 0 ; constructor ( protected items ) { super ( ) ; } summary ( ) { return `Array of ${ this . items . length } items` ; } async next ( ) { if ( this . trav >= this . items . length ) { return { value : null , done : true } ; } const item = this . items [ this . trav ] ; this . trav ++ ; return { value : deepClone ( item ) , done : false } ; } } class FunctionCallIterator < T > extends LazyIterator < T > { constructor ( protected nextFn ) { super ( ) ; } summary ( ) { return `Function call` ; } async next ( ) { try { return this . nextFn ( ) ; } catch ( e ) { e . message = `Error thrown while iterating through a dataset: ${ e . message }` ; throw e ; } } } class SerialIterator < T > extends LazyIterator < T > { private lastRead : Promise < IteratorResult < T > > ; constructor ( protected upstream ) { super ( ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Serial` ; } async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } private async serialNext ( ) { return this . upstream . next ( ) ; } } class SkipIterator < T > extends LazyIterator < T > { private lastRead : Promise < IteratorResult < T > > ; count = 0 ; constructor ( protected upstream , protected maxCount ) { super ( ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Skip` ; } async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } private async serialNext ( ) { while ( this . count ++ < this . maxCount ) { const skipped = await this . upstream . next ( ) ; if ( skipped . done ) { return skipped ; } tf . dispose ( skipped . value as { } ) ; } return this . upstream . next ( ) ; } } class TakeIterator < T > extends LazyIterator < T > { count = 0 ; constructor ( protected upstream , protected maxCount ) { super ( ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Take` ; } async next ( ) { if ( this . count ++ >= this . maxCount ) { return { value : null , done : true } ; } return this . upstream . next ( ) ; } } class RowMajorBatchIterator < T > extends LazyIterator < T [ ] > { private lastRead : Promise < IteratorResult < T [ ] > > ; constructor ( protected upstream , protected batchSize , protected enableSmallLastBatch = true ) { super ( ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> RowMajorBatch` ; } async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } private async serialNext ( ) { const batch = [ ] ; while ( batch . length < this . batchSize ) { const item = await this . upstream . next ( ) ; if ( item . done ) { if ( this . enableSmallLastBatch && batch . length > 0 ) { return { value : batch , done : false } ; } return { value : null , done : true } ; } batch . push ( item . value ) ; } return { value : batch , done : false } ; } } class FilterIterator < T > extends LazyIterator < T > { private lastRead : Promise < IteratorResult < T > > ; constructor ( protected upstream , protected predicate ) { super ( ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Filter` ; } async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } private async serialNext ( ) { while ( true ) { const item = await this . upstream . next ( ) ; if ( item . done || this . predicate ( item . value ) ) { return item ; } tf . dispose ( item . value as { } ) ; } } } class MapIterator < I , O > extends LazyIterator < O > { constructor ( protected upstream , protected transform ) { super ( ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Map` ; } async next ( ) { const item = await this . upstream . next ( ) ; if ( item . done ) { return { value : null , done : true } ; } const inputTensors = tf . tensor_util . getTensorsInContainer ( item . value as { } ) ; const mapped = this . transform ( item . value ) ; const outputTensors = tf . tensor_util . getTensorsInContainer ( mapped as { } ) ; for ( const t of inputTensors ) { if ( ! tf . tensor_util . isTensorInList ( t , outputTensors ) ) { t . dispose ( ) ; } } return { value : mapped , done : false } ; } } class ErrorHandlingLazyIterator < T > extends LazyIterator < T > { count = 0 ; constructor ( protected upstream , protected handler ) { super ( ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> handleErrors` ; } private lastRead : Promise < IteratorResult < T > > ; async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } async serialNext ( ) { while ( true ) { try { return await this . upstream . next ( ) ; } catch ( e ) { if ( ! this . handler ( e ) ) { return { value : null , done : true } ; } } } } } class AsyncMapIterator < I , O > extends LazyIterator < O > { constructor ( protected upstream , protected transform ) { super ( ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> AsyncMap` ; } async next ( ) { const item = await this . upstream . next ( ) ; if ( item . done ) { return { value : null , done : true } ; } const inputTensors = tf . tensor_util . getTensorsInContainer ( item . value as { } ) ; const mapped = await this . transform ( item . value ) ; const outputTensors = tf . tensor_util . getTensorsInContainer ( mapped as { } ) ; for ( const t of inputTensors ) { if ( ! tf . tensor_util . isTensorInList ( t , outputTensors ) ) { t . dispose ( ) ; } } return { value : mapped , done : false } ; } } /**
 * A base class for transforming streams that operate by maintaining an
 * output queue of elements that are ready to return via next().  This is
 * commonly required when the transformation is 1-to-many:  A call to next()
 * may trigger a call to the underlying stream, which will produce many
 * mapped elements of this stream-- of which we need to return only one, so
 * we have to queue the rest.
 */ export abstract class OneToManyIterator < T > extends LazyIterator < T > { private lastRead : Promise < IteratorResult < T > > ; protected outputQueue : RingBuffer < T > ; constructor ( ) { super ( ) ; this . outputQueue = new GrowingRingBuffer < T > ( ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } return protected abstract async pump ( ) ; async serialNext ( ) { while ( this . outputQueue . length ( ) === 0 ) { if ( ! await this . pump ( ) ) { return { value : null , done : true } ; } } return { value : this . outputQueue . shift ( ) , done : false } ; } } class FlatmapIterator < I , O > extends OneToManyIterator < O > { constructor ( protected upstream , protected transform ) { super ( ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Flatmap` ; } async pump ( ) { const item = await this . upstream . next ( ) ; if ( item . done ) { return false ; } const inputTensors = tf . tensor_util . getTensorsInContainer ( item . value as { } ) ; const mappedArray = this . transform ( item . value ) ; const outputTensors = tf . tensor_util . getTensorsInContainer ( mappedArray as { } ) ; this . outputQueue . pushAll ( mappedArray ) ; for ( const t of inputTensors ) { if ( ! tf . tensor_util . isTensorInList ( t , outputTensors ) ) { t . dispose ( ) ; } } return true ; } } /**
 * Provides a `LazyIterator` that concatenates a stream of underlying
 * streams.
 *
 * Doing this in a concurrency-safe way requires some trickery.  In
 * particular, we want this stream to return the elements from the
 * underlying streams in the correct order according to when next() was
 * called, even if the resulting Promises resolve in a different order.
 */ export class ChainedIterator < T > extends LazyIterator < T > { private lastRead : Promise < IteratorResult < T > > = null ; private iterator : LazyIterator < T > = null ; private moreIterators : LazyIterator < LazyIterator < T > > ; constructor ( iterators , private readonly baseErrorHandler ? ) { super ( ) ; this . moreIterators = iterators ; } summary ( ) { const upstreamSummaries = 'TODO: fill in upstream of chained summaries' ; return `${ upstreamSummaries } -> Chained` ; } async next ( ) { this . lastRead = this . readFromChain ( this . lastRead ) ; return this . lastRead ; } private async readFromChain ( lastRead ) { await lastRead ; if ( this . iterator == null ) { const iteratorResult = await this . moreIterators . next ( ) ; if ( iteratorResult . done ) { return { value : null , done : true } ; } this . iterator = iteratorResult . value ; if ( this . baseErrorHandler != null ) { this . iterator = this . iterator . handleErrors ( this . baseErrorHandler ) ; } } const itemResult = await this . iterator . next ( ) ; if ( itemResult . done ) { this . iterator = null ; return this . readFromChain ( lastRead ) ; } return itemResult ; } } export enum ZipMismatchMode { FAIL , SHORTEST , LONGEST } param iterators param mismatchMode class ZipIterator < O extends tf . TensorContainer > extends LazyIterator < O > { private count = 0 ; private currentPromise : Promise < IteratorResult < O > > = null ; constructor ( protected readonly iterators , protected readonly mismatchMode = ZipMismatchMode . FAIL ) { super ( ) ; } summary ( ) { const upstreamSummaries = 'TODO: fill in upstream of zip summaries' ; return `{${ upstreamSummaries }} -> Zip` ; } private async nextState ( afterState ) { await afterState ; let numIterators = 0 ; let iteratorsDone = 0 ; function getNext ( container ) { if ( container instanceof LazyIterator ) { const result = container . next ( ) ; return { value : result . then ( x => { numIterators ++ ; if ( x . done ) { iteratorsDone ++ ; } return x . value ; } ) , recurse : false } ; } else { return { value : null , recurse : true } ; } } const mapped = await deepMapAndAwaitAll ( this . iterators , getNext ) ; if ( numIterators === iteratorsDone ) { return { value : null , done : true } ; } if ( iteratorsDone > 0 ) { switch ( this . mismatchMode ) { case ZipMismatchMode . FAIL : throw new Error ( 'Zipped streams should have the same length. ' + `Mismatched at element ${ this . count }.` ) ; case ZipMismatchMode . SHORTEST : return { value : null , done : true } ; case ZipMismatchMode . LONGEST : default : } } this . count ++ ; return { value : mapped , done : false } ; } async next ( ) { this . currentPromise = this . nextState ( this . currentPromise ) ; return this . currentPromise ; } } /**
 * A stream that prefetches a given number of items from an upstream source,
 * returning them in FIFO order.
 *
 * Note this prefetches Promises, but makes no guarantees about when those
 * Promises resolve.
 */ export class PrefetchIterator < T > extends LazyIterator < T > { protected buffer : RingBuffer < Promise < IteratorResult < T > > > ; constructor ( protected upstream , protected bufferSize ) { super ( ) ; this . buffer = new RingBuffer < Promise < IteratorResult < T > > > ( bufferSize ) ; } summary ( ) { return `${ this . upstream . summary ( ) } -> Prefetch` ; } /**
     * Refill the prefetch buffer.  Returns only after the buffer is full, or
     * the upstream source is exhausted.
     */ protected refill ( ) { while ( ! this . buffer . isFull ( ) ) { const v = this . upstream . next ( ) ; this . buffer . push ( v ) ; } } next ( ) { this . refill ( ) ; return this . buffer . shift ( ) ; } } /**
 * A stream that performs a sliding-window random shuffle on an upstream
 * source. This is like a `PrefetchIterator` except that the items are
 * returned in randomized order.  Mixing naturally improves as the buffer
 * size increases.
 */ export class ShuffleIterator < T > extends PrefetchIterator < T > { private readonly random : seedrandom . prng ; private lastRead : Promise < IteratorResult < T > > ; private upstreamExhausted = false ; constructor ( protected upstream , protected windowSize , seed ? ) { super ( upstream , windowSize ) ; this . random = seedrandom . alea ( seed || tf . util . now ( ) . toString ( ) ) ; this . lastRead = Promise . resolve ( { value : null , done : false } ) ; } async next ( ) { this . lastRead = this . lastRead . then ( ( ) => this . serialNext ( ) ) ; return this . lastRead ; } private randomInt ( max ) { return Math . floor ( this . random ( ) * max ) ; } protected chooseIndex ( ) { return this . randomInt ( this . buffer . length ( ) ) ; } async serialNext ( ) { if ( ! this . upstreamExhausted ) { this . refill ( ) ; } while ( ! this . buffer . isEmpty ( ) ) { const chosenIndex = this . chooseIndex ( ) ; const result = await this . buffer . shuffleExcise ( chosenIndex ) ; if ( result . done ) { this . upstreamExhausted = true ; } else { this . refill ( ) ; return result ; } } return { value : null , done : true } ; } }