license import { AddN , AddNInputs , env , KernelConfig , KernelFunc , TensorInfo , upcastType } from '@tensorflow/tfjs-core' ; import { AddNProgram } from '../addn_gpu' ; import { AddNPackedProgram } from '../addn_packed_gpu' ; import { MathBackendWebGL } from '../backend_webgl' ; import { identity } from './Identity' ; export function addN ( args ) { const { inputs , backend } = args ; const tensors = inputs ; if ( tensors . length === 1 ) { return identity ( { inputs : { x : tensors [ 0 ] } , backend } ) ; } if ( tensors . length > env ( ) . get ( 'WEBGL_MAX_TEXTURES_IN_SHADER' ) ) { const midIndex = Math . floor ( tensors . length / 2 ) ; const leftSide = addN ( { inputs : tensors . slice ( 0 , midIndex ) , backend } ) ; const rightSide = addN ( { inputs : tensors . slice ( midIndex ) , backend } ) ; return addN ( { inputs : [ leftSide , rightSide ] , backend } ) ; } const dtype = tensors . map ( t => t . dtype ) . reduce ( ( d1 , d2 ) => upcastType ( d1 , d2 ) ) ; const shapes = tensors . map ( t => t . shape ) ; const usePackedOp = env ( ) . getBool ( 'WEBGL_PACK' ) ; const program = usePackedOp ? new AddNPackedProgram ( tensors [ 0 ] . shape , shapes ) : new AddNProgram ( tensors [ 0 ] . shape , shapes ) ; return backend . runWebGLProgram ( program , tensors , dtype ) ; } export const addNConfig = { kernelName : AddN , backendName : 'webgl' , kernelFunc : addN as { } as KernelFunc } ;