MLX
Loading...
Searching...
No Matches
Functions
softmax.h File Reference

Go to the source code of this file.

Functions

template<typename T >
softmax_exp (T x)
 
template<typename T , typename AccT = T, int N_READS = SOFTMAX_N_READS>
void softmax_single_row (const device T *in, device T *out, constant int &axis_size, uint gid, uint _lid, uint simd_lane_id, uint simd_group_id)
 
template<typename T , typename AccT = T, int N_READS = SOFTMAX_N_READS>
void softmax_looped (const device T *in, device T *out, constant int &axis_size, uint gid, uint lid, uint lsize, uint simd_lane_id, uint simd_group_id)
 

Function Documentation

◆ softmax_exp()

template<typename T >
T softmax_exp ( T x)
inline

◆ softmax_looped()

template<typename T , typename AccT = T, int N_READS = SOFTMAX_N_READS>
void softmax_looped ( const device T * in,
device T * out,
constant int & axis_size,
uint gid,
uint lid,
uint lsize,
uint simd_lane_id,
uint simd_group_id )

◆ softmax_single_row()

template<typename T , typename AccT = T, int N_READS = SOFTMAX_N_READS>
void softmax_single_row ( const device T * in,
device T * out,
constant int & axis_size,
uint gid,
uint _lid,
uint simd_lane_id,
uint simd_group_id )