license import { Diag , DiagInputs , KernelConfig , KernelFunc , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { DiagProgram } from '../diag_gpu' ; import { reshape } from './Reshape' ; export function diag ( args ) { const { inputs , backend } = args ; const { x } = inputs ; const outShape = [ ... x . shape , ... x . shape ] ; const xSize = util . sizeFromShape ( x . shape ) ; const flat = reshape ( { inputs : { x } , backend , attrs : { shape : [ xSize ] } } ) ; const program = new DiagProgram ( xSize ) ; const res = backend . runWebGLProgram ( program , [ flat ] , flat . dtype ) ; const out = reshape ( { inputs : { x : res } , backend , attrs : { shape : outShape } } ) ; backend . disposeIntermediateTensorInfo ( flat ) ; backend . disposeIntermediateTensorInfo ( res ) ; return out ; } export const diagConfig = { kernelName : Diag , backendName : 'webgl' , kernelFunc : diag as { } as KernelFunc } ;