license import { customGrad } from '../gradients' ; import { Tensor } from '../tensor' ; import { GradSaveFunc } from '../tensor_types' ; import { convertToTensor } from '../tensor_util_env' ; import { TensorLike } from '../types' ; import { cast } from './cast' ; import { exp } from './exp' ; import { log } from './log' ; import { max } from './max' ; import { mul } from './mul' ; import { op } from './operation' ; import { sub } from './sub' ; import { sum } from './sum' ; param logits param axis doc function logSoftmax_ < T extends Tensor > ( logits , axis = - 1 ) { const $logits = convertToTensor ( logits , 'logits' , 'logSoftmax' ) ; if ( axis === - 1 ) { axis = $logits . rank - 1 ; } if ( axis !== $logits . rank - 1 ) { throw Error ( 'Log Softmax along a non-last dimension is not yet supported. ' + `Logits was rank ${ $logits . rank } and axis was ${ axis }` ) ; } const customOp = customGrad ( ( logits , save ) => { const keepDims = true ; const xMax = max ( logits , axis , true ) ; const shifted = sub ( logits , xMax ) ; const value = sub ( cast ( shifted , 'float32' ) , log ( sum ( exp ( shifted ) , axis , keepDims ) ) ) ; save ( [ value ] ) ; const gradFunc = ( dy , saved ) => { const [ value ] = saved ; const keepDims = true ; const softmax = exp ( value ) ; return sub ( dy , mul ( sum ( dy , axis , keepDims ) , softmax ) ) ; } ; return { value , gradFunc } ; } ) ; return customOp ( $logits ) as T ; } export const logSoftmax = op ( { logSoftmax_ } ) ;