license import { KernelConfig , SparseReshape , SparseReshapeInputs , TensorInfo , TypedArray } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { sparseReshapeImplCPU } from '../kernel_utils/shared' ; export function sparseReshape ( args ) { const { inputs , backend } = args ; const { inputIndices , inputShape , newShape } = inputs ; if ( inputIndices . shape . length !== 2 ) { throw new Error ( `Input indices should be a matrix but received shape ${ inputIndices . shape }` ) ; } if ( inputShape . shape . length !== 1 ) { throw new Error ( `Input shape should be a vector but received shape ${ inputShape . shape }` ) ; } if ( newShape . shape . length !== 1 ) { throw new Error ( `Target shape should be a vector but received shape ${ newShape . shape }` ) ; } const $inputShape = Array . from ( backend . readSync ( inputShape . dataId ) as TypedArray ) ; const $inputIndices = backend . readSync ( inputIndices . dataId ) as TypedArray ; const targetShape = Array . from ( backend . readSync ( newShape . dataId ) as TypedArray ) ; const [ newIndices , indicesShape , outputShape ] = sparseReshapeImplCPU ( $inputIndices , inputIndices . shape , inputIndices . dtype , $inputShape , targetShape ) ; return [ backend . makeTensorInfo ( indicesShape , inputIndices . dtype , newIndices ) , backend . makeTensorInfo ( [ outputShape . length ] , newShape . dtype , new Int32Array ( outputShape ) ) , ] ; } export const sparseReshapeConfig = { kernelName : SparseReshape , backendName : 'webgl' , kernelFunc : sparseReshape , } ;