license import { memory , Tensor , tensor2d , tensor3d } from '@tensorflow/tfjs-core' ; import { test_util } from '@tensorflow/tfjs-core' ; import { TensorArray } from './tensor_array' ; let tensorArray ; let tensor ; let tensor2 ; const NAME = 'TA1' ; const DTYPE = 'int32' ; const SIZE = 10 ; const SHAPE = [ 1 , 1 ] ; const IDENTICAL_SHAPE = true ; const DYNAMIC_SIZE = false ; const CLEAR_AFTER_READ = false ; describe ( 'TensorArray' , ( ) => { beforeEach ( ( ) => { tensorArray = new TensorArray ( NAME , DTYPE , SIZE , SHAPE , IDENTICAL_SHAPE , DYNAMIC_SIZE , CLEAR_AFTER_READ ) ; tensor = tensor2d ( [ 1 ] , [ 1 , 1 ] , 'int32' ) ; tensor2 = tensor2d ( [ 2 ] , [ 1 , 1 ] , 'int32' ) ; } ) ; afterEach ( ( ) => tensorArray . clearAndClose ( ) ) ; it ( 'should initialize' , ( ) => { expect ( tensorArray . size ( ) ) . toEqual ( 0 ) ; expect ( tensorArray . name ) . toEqual ( 'TA1' ) ; expect ( tensorArray . dynamicSize ) . toBeFalsy ( ) ; expect ( tensorArray . closed ) . toBeFalsy ( ) ; } ) ; it ( 'should close' , ( ) => { const numOfTensors = memory ( ) . numTensors ; const size = tensorArray . size ( ) ; tensorArray . clearAndClose ( ) ; expect ( tensorArray . size ( ) ) . toBe ( 0 ) ; expect ( tensorArray . closed ) . toBeTruthy ( ) ; expect ( memory ( ) . numTensors ) . toEqual ( numOfTensors - size - 1 ) ; } ) ; it ( 'should not dispose keep tensors when close' , ( ) => { const numOfTensors = memory ( ) . numTensors ; tensorArray . write ( 0 , tensor ) ; tensorArray . write ( 1 , tensor2 ) ; const size = tensorArray . size ( ) ; const keepIds = new Set ( [ tensor . id ] ) ; tensorArray . clearAndClose ( keepIds ) ; expect ( tensorArray . size ( ) ) . toBe ( 0 ) ; expect ( tensorArray . closed ) . toBeTruthy ( ) ; expect ( tensor . isDisposed ) . toBeFalsy ( ) ; expect ( tensor2 . isDisposed ) . toBeTruthy ( ) ; expect ( memory ( ) . numTensors ) . toEqual ( numOfTensors - size ) ; } ) ; describe ( 'write' , ( ) => { it ( 'should add new tensor' , ( ) => { tensorArray . write ( 0 , tensor ) ; expect ( tensorArray . size ( ) ) . toBe ( 1 ) ; } ) ; it ( 'should add multiple tensors' , ( ) => { tensorArray . writeMany ( [ 0 , 1 ] , [ tensor , tensor2 ] ) ; expect ( tensorArray . size ( ) ) . toBe ( 2 ) ; } ) ; it ( 'should fail if dtype does not match' , ( ) => { const tensor = tensor2d ( [ 1 ] , [ 1 , 1 ] , 'float32' ) ; expect ( ( ) => tensorArray . write ( 0 , tensor ) ) . toThrow ( ) ; } ) ; it ( 'should fail if shape does not match' , ( ) => { const tensor = tensor2d ( [ 1 , 2 ] , [ 2 , 1 ] , 'int32' ) ; expect ( ( ) => tensorArray . write ( 0 , tensor ) ) . toThrow ( ) ; } ) ; it ( 'should allow unknown shape' , ( ) => { const unknownShape = undefined ; tensorArray = new TensorArray ( NAME , DTYPE , SIZE , unknownShape , IDENTICAL_SHAPE , DYNAMIC_SIZE , CLEAR_AFTER_READ ) ; const tensor = tensor2d ( [ 1 , 2 ] , [ 2 , 1 ] , 'int32' ) ; tensorArray . write ( 0 , tensor ) ; expect ( tensorArray . read ( 0 ) ) . not . toBeNull ( ) ; } ) ; it ( 'should fail if the index has already been written' , ( ) => { tensorArray . write ( 0 , tensor ) ; expect ( ( ) => tensorArray . write ( 0 , tensor ) ) . toThrow ( ) ; } ) ; it ( 'should fail if the index greater than array size' , ( ) => { expect ( ( ) => tensorArray . write ( 11 , tensor ) ) . toThrow ( ) ; } ) ; it ( 'should fail if the array is closed' , ( ) => { tensorArray . clearAndClose ( ) ; expect ( ( ) => tensorArray . write ( 0 , tensor ) ) . toThrow ( ) ; } ) ; it ( 'should create no new tensors' , ( ) => { const numTensors = memory ( ) . numTensors ; tensorArray . write ( 0 , tensor ) ; expect ( memory ( ) . numTensors ) . toEqual ( numTensors ) ; } ) ; } ) ; describe ( 'read' , ( ) => { beforeEach ( ( ) => { tensorArray . writeMany ( [ 0 , 1 ] , [ tensor , tensor2 ] ) ; } ) ; it ( 'should read the correct index' , ( ) => { expect ( tensorArray . read ( 0 ) ) . toBe ( tensor ) ; expect ( tensorArray . read ( 1 ) ) . toBe ( tensor2 ) ; } ) ; it ( 'should read the multiple indices' , ( ) => { expect ( tensorArray . readMany ( [ 0 , 1 ] ) ) . toEqual ( [ tensor , tensor2 ] ) ; } ) ; it ( 'should failed if index is out of bound' , ( ) => { expect ( ( ) => tensorArray . read ( 3 ) ) . toThrow ( ) ; expect ( ( ) => tensorArray . read ( - 1 ) ) . toThrow ( ) ; } ) ; it ( 'should failed if array is closed' , ( ) => { tensorArray . clearAndClose ( ) ; expect ( ( ) => tensorArray . read ( 0 ) ) . toThrow ( ) ; } ) ; it ( 'should create no new tensors' , ( ) => { const numTensors = memory ( ) . numTensors ; tensorArray . read ( 0 ) ; tensorArray . read ( 1 ) ; expect ( memory ( ) . numTensors ) . toEqual ( numTensors ) ; } ) ; } ) ; describe ( 'gather' , ( ) => { beforeEach ( ( ) => { tensorArray . writeMany ( [ 0 , 1 ] , [ tensor , tensor2 ] ) ; } ) ; it ( 'should return default packed tensors' , async ( ) => { const gathered = tensorArray . gather ( ) ; expect ( gathered . shape ) . toEqual ( [ 2 , 1 , 1 ] ) ; test_util . expectArraysClose ( await gathered . data ( ) , [ 1 , 2 ] ) ; } ) ; it ( 'should return packed tensors when indices is specified' , async ( ) => { const gathered = tensorArray . gather ( [ 1 , 0 ] ) ; expect ( gathered . shape ) . toEqual ( [ 2 , 1 , 1 ] ) ; test_util . expectArraysClose ( await gathered . data ( ) , [ 2 , 1 ] ) ; } ) ; it ( 'should return when indices longer than available tensors' , async ( ) => { const gathered = tensorArray . gather ( [ 1 , 0 , 2 , 3 ] ) ; expect ( gathered . shape ) . toEqual ( [ 2 , 1 , 1 ] ) ; test_util . expectArraysClose ( await gathered . data ( ) , [ 2 , 1 ] ) ; } ) ; it ( 'should fail if dtype is not matched' , ( ) => { expect ( ( ) => tensorArray . gather ( [ 0 , 1 ] , 'float32' ) ) . toThrow ( ) ; } ) ; it ( 'should create one new tensor' , ( ) => { const numTensors = memory ( ) . numTensors ; tensorArray . gather ( ) ; expect ( memory ( ) . numTensors ) . toEqual ( numTensors + 1 ) ; } ) ; } ) ; describe ( 'concat' , ( ) => { beforeEach ( ( ) => { tensorArray . writeMany ( [ 0 , 1 ] , [ tensor , tensor2 ] ) ; } ) ; it ( 'should return default concat tensors' , async ( ) => { const concat = tensorArray . concat ( ) ; expect ( concat . shape ) . toEqual ( [ 2 , 1 ] ) ; test_util . expectArraysClose ( await concat . data ( ) , [ 1 , 2 ] ) ; } ) ; it ( 'should fail if dtype is not matched' , ( ) => { expect ( ( ) => tensorArray . concat ( 'float32' ) ) . toThrow ( ) ; } ) ; it ( 'should create one new tensor' , ( ) => { const numTensors = memory ( ) . numTensors ; tensorArray . concat ( ) ; expect ( memory ( ) . numTensors ) . toEqual ( numTensors + 1 ) ; } ) ; } ) ; describe ( 'scatter' , ( ) => { it ( 'should scatter the input tensor' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; tensorArray . scatter ( [ 0 , 1 , 2 ] , input ) ; expect ( tensorArray . size ( ) ) . toEqual ( 3 ) ; } ) ; it ( 'should fail if indices and tensor shapes do not matched' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; expect ( ( ) => tensorArray . scatter ( [ 1 , 2 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should not fail if tensor array shape is -1' , ( ) => { tensorArray = new TensorArray ( NAME , DTYPE , SIZE , [ - 1 , 1 ] , IDENTICAL_SHAPE , DYNAMIC_SIZE , CLEAR_AFTER_READ ) ; const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; tensorArray . scatter ( [ 1 , 2 , 3 ] , input ) ; const res = tensorArray . gather ( [ 1 ] ) ; expect ( res ) . not . toBeNull ( ) ; } ) ; it ( 'should fail if max index > array max size' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; expect ( ( ) => tensorArray . scatter ( [ 1 , 2 , 11 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should fail if dtype is not matched' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'float32' ) ; expect ( ( ) => tensorArray . scatter ( [ 0 , 1 , 2 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should create three new tensors' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; const numTensors = memory ( ) . numTensors ; tensorArray . scatter ( [ 0 , 1 , 2 ] , input ) ; expect ( memory ( ) . numTensors ) . toEqual ( numTensors + 3 ) ; } ) ; } ) ; describe ( 'split' , ( ) => { beforeEach ( ( ) => { tensorArray = new TensorArray ( NAME , DTYPE , 3 , [ 2 , 2 ] , IDENTICAL_SHAPE , DYNAMIC_SIZE , CLEAR_AFTER_READ ) ; } ) ; it ( 'should split the input tensor' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] , [ 3 , 2 , 2 ] , 'int32' ) ; tensorArray . split ( [ 1 , 1 , 1 ] , input ) ; expect ( tensorArray . size ( ) ) . toEqual ( 3 ) ; } ) ; it ( 'should fail if indices and tensor shapes do not matched' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; expect ( ( ) => tensorArray . split ( [ 1 , 1 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should not fail if indices and tensor shapes is -1' , ( ) => { tensorArray = new TensorArray ( NAME , DTYPE , SIZE , [ - 1 , 1 ] , IDENTICAL_SHAPE , DYNAMIC_SIZE , CLEAR_AFTER_READ ) ; const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; expect ( ( ) => tensorArray . split ( [ 1 , 1 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should fail if length > array max size' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'int32' ) ; expect ( ( ) => tensorArray . split ( [ 1 , 1 , 1 , 1 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should fail if dtype is not matched' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 ] , [ 3 , 1 , 1 ] , 'float32' ) ; expect ( ( ) => tensorArray . split ( [ 1 , 1 , 1 ] , input ) ) . toThrow ( ) ; } ) ; it ( 'should create three new tensors' , ( ) => { const input = tensor3d ( [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] , [ 3 , 2 , 2 ] , 'int32' ) ; const numTensors = memory ( ) . numTensors ; tensorArray . split ( [ 1 , 1 , 1 ] , input ) ; expect ( memory ( ) . numTensors ) . toEqual ( numTensors + 3 ) ; } ) ; } ) ; } ) ;