use crate::ml::autograd::node::operation::BackwardStatus;
use crate::ml::autograd::tape::backward::accumulate_value_gradient;
use crate::ml::matrix::UpsampleMode;
use crate::{Matrix, Result};
use std::collections::HashMap;
pub(in crate::ml::autograd) struct GradNodeUpsample2d {
pub(in crate::ml::autograd) input: Matrix,
pub(in crate::ml::autograd) scale_factor: usize,
pub(in crate::ml::autograd) mode: UpsampleMode,
pub(in crate::ml::autograd) output_id: u64,
}
impl GradNodeUpsample2d {
pub(in crate::ml::autograd) fn backward(
&self,
gradients: &mut HashMap<u64, Matrix>,
) -> Result<BackwardStatus> {
let Self {
input,
output_id,
scale_factor,
mode,
} = self;
let Some(output_gradient) = gradients.remove(output_id) else {
return Ok(BackwardStatus::Skipped);
};
let input_gradient =
crate::ml::matrix::upsample_2d_backward(input, &output_gradient, *scale_factor, *mode)?;
accumulate_value_gradient(gradients, input.value_id(), input_gradient)?;
Ok(BackwardStatus::Recorded)
}
}