use crate::ml::autograd::node::operation::BackwardStatus;
use crate::ml::autograd::tape::backward::accumulate_value_gradient;
use crate::ml::matrix as ml_matrix;
use crate::{Matrix, Result};
use std::collections::HashMap;
pub(in crate::ml::autograd) struct GradNodeRope {
pub(in crate::ml::autograd) input: Matrix,
pub(in crate::ml::autograd) num_heads: usize,
pub(in crate::ml::autograd) head_dim: usize,
pub(in crate::ml::autograd) theta_base: f32,
pub(in crate::ml::autograd) position_offset: u32,
pub(in crate::ml::autograd) output_id: u64,
}
impl GradNodeRope {
pub(in crate::ml::autograd) fn backward(
&self,
gradients: &mut HashMap<u64, Matrix>,
) -> Result<BackwardStatus> {
let Self {
input,
output_id,
num_heads,
head_dim,
theta_base,
position_offset,
} = self;
let Some(output_gradient) = gradients.remove(output_id) else {
return Ok(BackwardStatus::Skipped);
};
let input_gradient = ml_matrix::rope_backward(
&output_gradient,
*num_heads,
*head_dim,
*theta_base,
*position_offset,
)?;
accumulate_value_gradient(gradients, input.value_id(), input_gradient)?;
Ok(BackwardStatus::Recorded)
}
}