license import { KernelFunc , Max , MaxAttrs , MaxInputs , TensorInfo } from '@tensorflow/tfjs-core' ; import { backend_util , KernelConfig , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { maxImplCPU } from '../kernel_utils/shared' ; import { maxImpl } from './Max_impl' ; import { transposeImpl , transposeImplCPU } from './Transpose_impl' ; export function max ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { reductionIndices , keepDims } = attrs ; const xRank = x . shape . length ; const origAxes = util . parseAxisParam ( reductionIndices , x . shape ) ; let axes = origAxes ; const permutedAxes = backend_util . getAxesPermutation ( axes , xRank ) ; const maxInputIsTransposed = permutedAxes != null ; const shouldExecuteOnCPU = backend . shouldExecuteOnCPU ( [ x ] ) ; let maxInput = x ; if ( maxInputIsTransposed ) { if ( shouldExecuteOnCPU ) { const xTexData = backend . texData . get ( maxInput . dataId ) ; const values = xTexData . values as TypedArray ; const newShape = new Array ( xRank ) ; for ( let i = 0 ; i < newShape . length ; i ++ ) { newShape [ i ] = x . shape [ permutedAxes [ i ] ] ; } const maxInputValues = transposeImplCPU ( values , x . shape , x . dtype , permutedAxes , newShape ) ; maxInput = backend . makeTensorInfo ( newShape , x . dtype ) ; const maxInputData = backend . texData . get ( maxInput . dataId ) ; maxInputData . values = maxInputValues ; } else { maxInput = transposeImpl ( x , permutedAxes , backend ) ; } axes = backend_util . getInnerMostAxes ( axes . length , xRank ) ; } backend_util . assertAxesAreInnerMostDims ( 'max' , axes , xRank ) ; const [ maxOutShape , reduceShape ] = backend_util . computeOutAndReduceShapes ( maxInput . shape , axes ) ; let outShape = maxOutShape ; if ( keepDims ) { outShape = backend_util . expandShapeToKeepDim ( maxOutShape , origAxes ) ; } let out ; if ( shouldExecuteOnCPU ) { const xTexData = backend . texData . get ( maxInput . dataId ) ; const values = xTexData . values as TypedArray ; const outValues = maxImplCPU ( values , util . sizeFromShape ( reduceShape ) , outShape , x . dtype ) ; out = backend . makeTensorInfo ( outShape , x . dtype ) ; const outData = backend . texData . get ( out . dataId ) ; outData . values = outValues ; } else { out = maxImpl ( maxInput , reduceShape , outShape , backend ) ; } if ( maxInputIsTransposed ) { backend . disposeIntermediateTensorInfo ( maxInput ) ; } return out ; } export const maxConfig = { kernelName : Max , backendName : 'webgl' , kernelFunc : max as { } as KernelFunc } ;