pub fn rprop_step<T: Triton, const BLOCK_SIZE: i32>(
params_ptr: T::Pointer<f32>,
grad_ptr: T::Pointer<f32>,
prev_grad_ptr: T::Pointer<f32>,
step_size_ptr: T::Pointer<f32>,
n_elements: i32,
eta_plus: f32,
eta_minus: f32,
step_min: f32,
step_max: f32,
)where
T::I32Tensor: Tensor<i32, 1> + Comparison<i32, BoolTensor = T::BoolTensor>,
T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,Expand description
Rprop step (resilient backpropagation).
sign = g * prev_g (product of current and previous gradient)
step_size = step_size * eta_plus if sign > 0
= step_size * eta_minus if sign < 0
= step_size otherwise
step_size = clamp(step_size, step_min, step_max)
g_masked = 0 if sign < 0 (gradient reversal: skip update)
= g otherwise
p -= sign(g_masked) * step_size
prev_g = g_maskedGrid: [ceil(n_elements / BLOCK_SIZE), 1, 1].