license import { concat , DataType , keep , reshape , scalar , slice , stack , Tensor , tensor , tidy , unstack } from '@tensorflow/tfjs-core' ; import { assertShapesMatchAllowUndefinedSize , inferElementShape , mergeElementShape } from './tensor_utils' ; /**
 * TensorList stores a container of `tf.Tensor` objects, which are accessible
 * via tensors field.
 *
 * In order to get a copy of the underlying list, use the copy method:
 * ```
 *    TensorList b = a.copy();
 *    b.tensors().pushBack(t);  // This does not modify a.tensors().
 * ```
 *
 * Note that this is not a deep copy: the memory locations of the underlying
 * tensors will still point to the same locations of the corresponding tensors
 * in the original.
 */ export class TensorList { readonly idTensor : Tensor ; maxNumElements : number ; get id ( ) { return this . idTensor . id ; } param tensors param elementShape param elementDtype param maxNumElements constructor ( readonly tensors , readonly elementShape , readonly elementDtype , maxNumElements = - 1 ) { if ( tensors != null ) { tensors . forEach ( tensor => { if ( elementDtype !== tensor . dtype ) { throw new Error ( `Invalid data types; op elements ${ elementDtype }, but list elements ${ tensor . dtype }` ) ; } assertShapesMatchAllowUndefinedSize ( elementShape , tensor . shape , 'TensorList shape mismatch: ' ) ; keep ( tensor ) ; } ) ; } this . idTensor = scalar ( 0 ) ; this . maxNumElements = maxNumElements ; keep ( this . idTensor ) ; } /**
     * Get a new TensorList containing a copy of the underlying tensor container.
     */ copy ( ) { return new TensorList ( [ ... this . tensors ] , this . elementShape , this . elementDtype ) ; } /**
     * Dispose the tensors and idTensor and clear the tensor list.
     */ clearAndClose ( keepIds ? ) { this . tensors . forEach ( tensor => { if ( keepIds == null || ! keepIds . has ( tensor . id ) ) { tensor . dispose ( ) ; } } ) ; this . tensors . length = 0 ; this . idTensor . dispose ( ) ; } /**
     * The size of the tensors in the tensor list.
     */ size ( ) { return this . tensors . length ; } param elementShape param elementDtype param numElements stack ( elementShape , elementDtype , numElements = - 1 ) { if ( elementDtype !== this . elementDtype ) { throw new Error ( `Invalid data types; op elements ${ elementDtype }, but list elements ${ this . elementDtype }` ) ; } if ( numElements !== - 1 && this . tensors . length !== numElements ) { throw new Error ( `Operation expected a list with ${ numElements } elements but got a list with ${ this . tensors . length } elements.` ) ; } assertShapesMatchAllowUndefinedSize ( elementShape , this . elementShape , 'TensorList shape mismatch: ' ) ; const outputElementShape = inferElementShape ( this . elementShape , this . tensors , elementShape ) ; return tidy ( ( ) => { const reshapedTensors = this . tensors . map ( tensor => reshape ( tensor , outputElementShape ) ) ; return stack ( reshapedTensors , 0 ) ; } ) ; } param elementShape param elementDtype popBack ( elementShape , elementDtype ) { if ( elementDtype !== this . elementDtype ) { throw new Error ( `Invalid data types; op elements ${ elementDtype }, but list elements ${ this . elementDtype }` ) ; } if ( this . size ( ) === 0 ) { throw new Error ( 'Trying to pop from an empty list.' ) ; } const outputElementShape = inferElementShape ( this . elementShape , this . tensors , elementShape ) ; const tensor = this . tensors . pop ( ) ; assertShapesMatchAllowUndefinedSize ( tensor . shape , elementShape , 'TensorList shape mismatch: ' ) ; return reshape ( tensor , outputElementShape ) ; } param tensor pushBack ( tensor ) { if ( tensor . dtype !== this . elementDtype ) { throw new Error ( `Invalid data types; op elements ${ tensor . dtype }, but list elements ${ this . elementDtype }` ) ; } assertShapesMatchAllowUndefinedSize ( tensor . shape , this . elementShape , 'TensorList shape mismatch: ' ) ; if ( this . maxNumElements === this . size ( ) ) { throw new Error ( `Trying to push element into a full list.` ) ; } keep ( tensor ) ; this . tensors . push ( tensor ) ; } param size resize ( size ) { if ( size < 0 ) { throw new Error ( `TensorListResize expects size to be non-negative. Got: ${ size }` ) ; } if ( this . maxNumElements !== - 1 && size > this . maxNumElements ) { throw new Error ( `TensorListResize input size ${ size } is greater maxNumElement ${ this . maxNumElements }.` ) ; } this . tensors . length = size ; } param elementShape param elementDtype param elementIndex getItem ( elementIndex , elementShape , elementDtype ) { if ( elementDtype !== this . elementDtype ) { throw new Error ( `Invalid data types; op elements ${ elementDtype }, but list elements ${ this . elementDtype }` ) ; } if ( elementIndex < 0 || elementIndex > this . tensors . length ) { throw new Error ( `Trying to access element ${ elementIndex } in a list with ${ this . tensors . length } elements.` ) ; } if ( this . tensors [ elementIndex ] == null ) { throw new Error ( `element at index ${ elementIndex } is null.` ) ; } assertShapesMatchAllowUndefinedSize ( this . tensors [ elementIndex ] . shape , elementShape , 'TensorList shape mismatch: ' ) ; const outputElementShape = inferElementShape ( this . elementShape , this . tensors , elementShape ) ; return reshape ( this . tensors [ elementIndex ] , outputElementShape ) ; } param elementIndex param tensor setItem ( elementIndex , tensor ) { if ( tensor . dtype !== this . elementDtype ) { throw new Error ( `Invalid data types; op elements ${ tensor . dtype }, but list elements ${ this . elementDtype }` ) ; } if ( elementIndex < 0 || this . maxNumElements !== - 1 && elementIndex >= this . maxNumElements ) { throw new Error ( `Trying to set element ${ elementIndex } in a list with max ${ this . maxNumElements } elements.` ) ; } assertShapesMatchAllowUndefinedSize ( this . elementShape , tensor . shape , 'TensorList shape mismatch: ' ) ; keep ( tensor ) ; this . tensors [ elementIndex ] = tensor ; } param indices param elementDtype param elementShape gather ( indices , elementDtype , elementShape ) { if ( elementDtype !== this . elementDtype ) { throw new Error ( `Invalid data types; op elements ${ elementDtype }, but list elements ${ this . elementDtype }` ) ; } assertShapesMatchAllowUndefinedSize ( this . elementShape , elementShape , 'TensorList shape mismatch: ' ) ; indices = indices . slice ( 0 , this . size ( ) ) ; const outputElementShape = inferElementShape ( this . elementShape , this . tensors , elementShape ) ; if ( indices . length === 0 ) { return tensor ( [ ] , [ 0 ] . concat ( outputElementShape ) ) ; } return tidy ( ( ) => { const tensors = indices . map ( i => reshape ( this . tensors [ i ] , outputElementShape ) ) ; return stack ( tensors , 0 ) ; } ) ; } param elementDtype param elementShape concat ( elementDtype , elementShape ) { if ( ! ! elementDtype && elementDtype !== this . elementDtype ) { throw new Error ( `TensorList dtype is ${ this . elementDtype } but concat requested dtype ${ elementDtype }` ) ; } assertShapesMatchAllowUndefinedSize ( this . elementShape , elementShape , 'TensorList shape mismatch: ' ) ; const outputElementShape = inferElementShape ( this . elementShape , this . tensors , elementShape ) ; if ( this . size ( ) === 0 ) { return tensor ( [ ] , [ 0 ] . concat ( outputElementShape ) ) ; } return tidy ( ( ) => { const tensors = this . tensors . map ( t => reshape ( t , outputElementShape ) ) ; return concat ( tensors , 0 ) ; } ) ; } } param tensor param elementShape export function fromTensor ( tensor , elementShape , elementDtype ) { const dtype = tensor . dtype ; if ( tensor . shape . length < 1 ) { throw new Error ( `Tensor must be at least a vector, but saw shape: ${ tensor . shape }` ) ; } if ( tensor . dtype !== elementDtype ) { throw new Error ( `Invalid data types; op elements ${ tensor . dtype }, but list elements ${ elementDtype }` ) ; } const tensorElementShape = tensor . shape . slice ( 1 ) ; assertShapesMatchAllowUndefinedSize ( tensorElementShape , elementShape , 'TensorList shape mismatch: ' ) ; const tensorList = unstack ( tensor ) ; return new TensorList ( tensorList , elementShape , dtype ) ; } param elementShape param elementDtype param numElements export function reserve ( elementShape , elementDtype , numElements ) { return new TensorList ( [ ] , elementShape , elementDtype , numElements ) ; } param indices param tensor param elementShape param numElements export function scatter ( tensor , indices , elementShape , numElements ? ) { if ( indices . length !== tensor . shape [ 0 ] ) { throw new Error ( `Expected len(indices) == tensor.shape[0], but saw: ${ indices . length } vs. ${ tensor . shape [ 0 ] }` ) ; } const maxIndex = Math . max ( ... indices ) ; if ( numElements != null && numElements !== - 1 && maxIndex >= numElements ) { throw new Error ( `Max index must be < array size (${ maxIndex }  vs. ${ numElements })` ) ; } const list = new TensorList ( [ ] , elementShape , tensor . dtype , numElements ) ; const tensors = unstack ( tensor , 0 ) ; indices . forEach ( ( value , index ) => { list . setItem ( value , tensors [ index ] ) ; } ) ; return list ; } param length param tensor param elementShape export function split ( tensor , length , elementShape ) { let totalLength = 0 ; const cumulativeLengths = length . map ( len => { totalLength += len ; return totalLength ; } ) ; if ( totalLength !== tensor . shape [ 0 ] ) { throw new Error ( `Expected sum of lengths to be equal to
          tensor.shape[0], but sum of lengths is
        ${ totalLength }, and tensor's shape is: ${ tensor . shape }` ) ; } const shapeWithoutFirstDim = tensor . shape . slice ( 1 ) ; const outputElementShape = mergeElementShape ( shapeWithoutFirstDim , elementShape ) ; const elementPerRow = totalLength === 0 ? 0 : tensor . size / totalLength ; const tensors = tidy ( ( ) => { const tensors = [ ] ; tensor = reshape ( tensor , [ 1 , totalLength , elementPerRow ] ) ; for ( let i = 0 ; i < length . length ; ++ i ) { const previousLength = ( i === 0 ) ? 0 : cumulativeLengths [ i - 1 ] ; const indices = [ 0 , previousLength , 0 ] ; const sizes = [ 1 , length [ i ] , elementPerRow ] ; tensors [ i ] = reshape ( slice ( tensor , indices , sizes ) , outputElementShape as number [ ] ) ; } tensor . dispose ( ) ; return tensors ; } ) ; const list = new TensorList ( [ ] , elementShape , tensor . dtype , length . length ) ; for ( let i = 0 ; i < tensors . length ; i ++ ) { list . setItem ( i , tensors [ i ] ) ; } return list ; }