license import { _FusedMatMul , _FusedMatMulAttrs , _FusedMatMulInputs , KernelConfig , KernelFunc , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { applyActivation } from '../utils/fused_utils' ; import { add } from './Add' ; import { batchMatMul } from './BatchMatMul' ; export function _fusedMatMul ( args ) { const { inputs , backend , attrs } = args ; const { a , b , bias , preluActivationWeights } = inputs ; const { transposeA , transposeB , activation , leakyreluAlpha } = attrs ; let current ; let addRes ; let activationRes ; const intermediates = [ ] ; const matMulRes = batchMatMul ( { inputs : { a , b } , attrs : { transposeA , transposeB } , backend } ) ; current = matMulRes ; if ( bias ) { addRes = add ( { inputs : { a : current , b : bias } , backend } ) as TensorInfo ; intermediates . push ( current ) ; current = addRes ; } if ( activation ) { activationRes = applyActivation ( backend , current , activation , preluActivationWeights , leakyreluAlpha ) ; intermediates . push ( current ) ; current = activationRes ; } for ( const i of intermediates ) { backend . disposeIntermediateTensorInfo ( i ) ; } return current ; } export const _fusedMatMulConfig = { kernelName : _FusedMatMul , backendName : 'cpu' , kernelFunc : _fusedMatMul as { } as KernelFunc , } ;