license import { DepthToSpace , DepthToSpaceAttrs , DepthToSpaceInputs , KernelConfig , KernelFunc , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; export function depthToSpace ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { blockSize , dataFormat } = attrs ; util . assert ( dataFormat === 'NHWC' , ( ) => `Only NHWC dataFormat supported on CPU for depthToSpace. Got ${ dataFormat }` ) ; const batchSize = x . shape [ 0 ] ; const inputHeight = x . shape [ 1 ] ; const inputWidth = x . shape [ 2 ] ; const inputDepth = x . shape [ 3 ] ; const outputHeight = inputHeight * blockSize ; const outputWidth = inputWidth * blockSize ; const outputDepth = inputDepth / ( blockSize * blockSize ) ; const xValues = backend . data . get ( x . dataId ) . values as TypedArray ; const result = new Float32Array ( batchSize * outputHeight * outputWidth * outputDepth ) ; let outputIdx = 0 ; for ( let b = 0 ; b < batchSize ; ++ b ) { for ( let h = 0 ; h < outputHeight ; ++ h ) { const inH = Math . floor ( h / blockSize ) ; const offsetH = ( h % blockSize ) ; for ( let w = 0 ; w < outputWidth ; ++ w ) { const inW = Math . floor ( w / blockSize ) ; const offsetW = ( w % blockSize ) ; const offsetD = ( offsetH * blockSize + offsetW ) * outputDepth ; for ( let d = 0 ; d < outputDepth ; ++ d ) { const inD = d + offsetD ; const inputIdx = inD + inputDepth * ( inW + inputWidth * ( inH + inputHeight * b ) ) ; result [ outputIdx ++ ] = xValues [ inputIdx ] ; } } } } return backend . makeTensorInfo ( [ batchSize , outputHeight , outputWidth , outputDepth ] , x . dtype , result ) ; } export const depthToSpaceConfig = { kernelName : DepthToSpace , backendName : 'cpu' , kernelFunc : depthToSpace as { } as KernelFunc } ;