oarust 0.8.6

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=ml_training.json schema_version=1 generator_version=64
// schema_sha256=6678961ee57f24310158676c7b4fe0821e21ca48db7edcb20d2a749a164976ae
// Saved-state and backward recipes; lifecycle remains tape-owned.
use crate::ml::autograd::node::operation::BackwardStatus;
use crate::ml::autograd::tape::backward::accumulate_value_gradient;
use crate::{Matrix, Result};
use crate::{matrix, ml::loss};
use std::collections::HashMap;

pub(in crate::ml::autograd) struct GradNodePpoClippedPolicy {
	pub(in crate::ml::autograd) new_log_probability: Matrix,
	pub(in crate::ml::autograd) old_log_probability: Matrix,
	pub(in crate::ml::autograd) advantage: Matrix,
	pub(in crate::ml::autograd) clip_epsilon: f32,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodePpoClippedPolicy {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			new_log_probability,
			old_log_probability,
			advantage,
			clip_epsilon,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(
			&loss::ppo_clipped_policy_backward(
				new_log_probability,
				old_log_probability,
				advantage,
				*clip_epsilon,
			)?,
			&output_gradient,
		)?;
		accumulate_value_gradient(gradients, new_log_probability.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeCrossEntropy {
	pub(in crate::ml::autograd) logits: Matrix,
	pub(in crate::ml::autograd) targets: Matrix,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeCrossEntropy {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			logits,
			targets,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(
			&loss::cross_entropy_backward(logits, targets)?,
			&output_gradient,
		)?;
		accumulate_value_gradient(gradients, logits.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeMaskedCrossEntropy {
	pub(in crate::ml::autograd) logits: Matrix,
	pub(in crate::ml::autograd) targets: Matrix,
	pub(in crate::ml::autograd) mask: Matrix,
	pub(in crate::ml::autograd) valid_count: usize,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeMaskedCrossEntropy {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			logits,
			targets,
			mask,
			valid_count,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(
			&loss::cross_entropy_masked_backward(logits, targets, mask, *valid_count)?,
			&output_gradient,
		)?;
		accumulate_value_gradient(gradients, logits.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeSmoothL1 {
	pub(in crate::ml::autograd) prediction: Matrix,
	pub(in crate::ml::autograd) target: Matrix,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeSmoothL1 {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			prediction,
			target,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(
			&loss::smooth_l1_backward(prediction, target)?,
			&output_gradient,
		)?;
		accumulate_value_gradient(gradients, prediction.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeMse {
	pub(in crate::ml::autograd) prediction: Matrix,
	pub(in crate::ml::autograd) target: Matrix,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeMse {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			prediction,
			target,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(&loss::mse_backward(prediction, target)?, &output_gradient)?;
		accumulate_value_gradient(gradients, prediction.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeL1 {
	pub(in crate::ml::autograd) prediction: Matrix,
	pub(in crate::ml::autograd) target: Matrix,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeL1 {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			prediction,
			target,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(&loss::l1_backward(prediction, target)?, &output_gradient)?;
		accumulate_value_gradient(gradients, prediction.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeBce {
	pub(in crate::ml::autograd) prediction: Matrix,
	pub(in crate::ml::autograd) target: Matrix,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeBce {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			prediction,
			target,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let gradient = matrix::mul(&loss::bce_backward(prediction, target)?, &output_gradient)?;
		accumulate_value_gradient(gradients, prediction.value_id(), gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}