license import { KernelFunc , Max , MaxAttrs , MaxInputs , TensorInfo } from '@tensorflow/tfjs-core' ; import { backend_util , KernelConfig } from '@tensorflow/tfjs-core' ; import { TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; import { maxImpl } from './Max_impl' ; import { transposeImpl } from './Transpose_impl' ; export function max ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { reductionIndices , keepDims } = attrs ; const cpuBackend = backend ; let xShape = x . shape ; const xRank = xShape . length ; const origAxes = util . parseAxisParam ( reductionIndices , xShape ) ; let axes = origAxes ; const permutedAxes = backend_util . getAxesPermutation ( axes , xRank ) ; let xVals = cpuBackend . data . get ( x . dataId ) . values as TypedArray ; if ( permutedAxes != null ) { const newShape = new Array ( xRank ) ; for ( let i = 0 ; i < newShape . length ; i ++ ) { newShape [ i ] = xShape [ permutedAxes [ i ] ] ; } xVals = transposeImpl ( xVals , xShape , x . dtype , permutedAxes , newShape ) ; axes = backend_util . getInnerMostAxes ( axes . length , xRank ) ; xShape = newShape ; } assertNotComplex ( x , 'max' ) ; backend_util . assertAxesAreInnerMostDims ( 'max' , axes , xRank ) ; const [ maxOutShape , reduceShape ] = backend_util . computeOutAndReduceShapes ( xShape , axes ) ; const reduceSize = util . sizeFromShape ( reduceShape ) ; const result = maxImpl ( xVals , reduceSize , maxOutShape , x . dtype ) ; const dataId = cpuBackend . write ( result , maxOutShape , x . dtype ) ; let outShape = maxOutShape ; if ( keepDims ) { const newShape = backend_util . expandShapeToKeepDim ( maxOutShape , origAxes ) ; outShape = newShape ; } return { dataId , shape : outShape , dtype : x . dtype } ; } export const maxConfig = { kernelName : Max , backendName : 'cpu' , kernelFunc : max as { } as KernelFunc } ;