license import { buffer , Diag , DiagInputs , KernelConfig , KernelFunc , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; export function diag ( args ) { const { inputs , backend } = args ; const { x } = inputs ; const xSize = util . sizeFromShape ( x . shape ) ; const xVals = backend . data . get ( x . dataId ) . values as TypedArray ; const outBuf = buffer ( [ xSize , xSize ] , x . dtype ) ; const vals = outBuf . values ; for ( let i = 0 ; i < xVals . length ; i ++ ) { vals [ i * xSize + i ] = xVals [ i ] ; } const outShape = [ ... x . shape , ... x . shape ] ; return backend . makeTensorInfo ( outShape , outBuf . dtype , outBuf . values ) ; } export const diagConfig = { kernelName : Diag , backendName : 'cpu' , kernelFunc : diag as { } as KernelFunc } ;