license import { backend_util , Cumsum , CumsumAttrs , CumsumInputs , KernelConfig , KernelFunc , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { CumSumProgram } from '../cumsum_gpu' ; import { identity } from './Identity' ; import { transpose } from './Transpose' ; export function cumsum ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { axis , exclusive , reverse } = attrs ; const xRank = x . shape . length ; const permutation = backend_util . getAxesPermutation ( [ axis ] , xRank ) ; let permutedX = x ; if ( permutation != null ) { permutedX = transpose ( { inputs : { x } , backend , attrs : { perm : permutation } } ) ; } const permutedAxis = backend_util . getInnerMostAxes ( 1 , xRank ) [ 0 ] ; if ( permutedAxis !== xRank - 1 ) { throw new Error ( `WebGL cumsum shader expects an inner-most axis=${ x . shape . length - 1 } ` + `but got axis=${ axis }` ) ; } const size = permutedX . shape [ permutedAxis ] ; let result = identity ( { inputs : { x : permutedX } , backend } ) ; for ( let i = 0 ; i <= Math . ceil ( Math . log2 ( size ) ) - 1 ; i ++ ) { const program = new CumSumProgram ( permutedX . shape , false , reverse ) ; const customValues = [ [ i ] ] ; const prevResult = result ; result = backend . runWebGLProgram ( program , [ result ] , result . dtype , customValues ) ; backend . disposeIntermediateTensorInfo ( prevResult ) ; } if ( exclusive ) { const program = new CumSumProgram ( permutedX . shape , exclusive , reverse ) ; const prevResult = result ; result = backend . runWebGLProgram ( program , [ result ] , result . dtype ) ; backend . disposeIntermediateTensorInfo ( prevResult ) ; } if ( permutation != null ) { const reversePermutation = backend_util . getUndoAxesPermutation ( permutation ) ; const reverseTransposedResult = transpose ( { inputs : { x : result } , backend , attrs : { perm : reversePermutation } } ) ; backend . disposeIntermediateTensorInfo ( result ) ; backend . disposeIntermediateTensorInfo ( permutedX ) ; return reverseTransposedResult ; } return result ; } export const cumsumConfig = { kernelName : Cumsum , backendName : 'webgl' , kernelFunc : cumsum as { } as KernelFunc } ;