license import { backend_util , Einsum , EinsumAttrs , EinsumInputs , KernelConfig , KernelFunc , Tensor , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { multiply } from './Multiply' ; import { reshape } from './Reshape' ; import { sum } from './Sum' ; import { transpose } from './Transpose' ; export function einsum ( args ) { const { inputs , backend , attrs } = args ; const { equation } = attrs ; const tensors = inputs as Tensor [ ] ; const { allDims , summedDims , idDims } = backend_util . decodeEinsumEquation ( equation , tensors . length ) ; backend_util . checkEinsumDimSizes ( allDims . length , idDims , tensors ) ; const { path , steps } = backend_util . getEinsumComputePath ( summedDims , idDims ) ; const nSteps = steps . length ; let out = null ; let numDimsRemaining = allDims . length ; const tensorsToDispose = [ ] ; for ( let i = 0 ; i < nSteps ; ++ i ) { for ( const idTerm of steps [ i ] ) { const { permutationIndices : perm , expandDims : dimsToExpand } = backend_util . getEinsumPermutation ( numDimsRemaining , idDims [ idTerm ] ) ; let x ; if ( backend_util . isIdentityPermutation ( perm ) ) { x = tensors [ idTerm ] ; } else { x = transpose ( { inputs : { x : tensors [ idTerm ] } , backend , attrs : { perm } } ) ; tensorsToDispose . push ( x ) ; } const targetShape = x . shape . slice ( ) ; for ( let k = 0 ; k < dimsToExpand . length ; ++ k ) { targetShape . splice ( dimsToExpand [ k ] , 0 , 1 ) ; } if ( ! util . arraysEqual ( x . shape , targetShape ) ) { x = reshape ( { inputs : { x } , backend , attrs : { shape : targetShape } } ) ; tensorsToDispose . push ( x ) ; } if ( out === null ) { out = x ; } else { out = multiply ( { inputs : { a : x , b : out } , backend } ) as TensorInfo ; tensorsToDispose . push ( out ) ; } } if ( i < nSteps - 1 ) { if ( path [ i ] >= 0 ) { out = sum ( { inputs : { x : out } , backend , attrs : { axis : path [ i ] - ( allDims . length - numDimsRemaining ) , keepDims : false } } ) ; tensorsToDispose . push ( out ) ; } numDimsRemaining -- ; } } for ( const tensorInfo of tensorsToDispose ) { if ( tensorInfo === out ) { continue ; } backend . disposeIntermediateTensorInfo ( tensorInfo ) ; } return out ; } export const einsumConfig = { kernelName : Einsum , backendName : 'webgl' , kernelFunc : einsum as { } as KernelFunc } ;