corpus_ipc/models/
gradients.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use serde::{Deserialize, Serialize};
4
5use crate::validation::{
6    ProtocolLimits, Validate, ValidationError, add_to_total, check_count, check_finite_slice,
7    check_string, check_unique_strings,
8};
9
10use super::de::{
11    de_eligibility_trace, de_gradient_rows, de_gradient_values, de_layer_id, de_session_id,
12    prefix_path,
13};
14
15/// Gradient update batch from external training or optimization algorithms.
16#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
17#[serde(try_from = "GradientBatchWire")]
18pub struct GradientBatch {
19    /// Session ID for routing back to correct experiment.
20    pub session_id: String,
21    /// Batch ID correlation with original input.
22    pub batch_id: u64,
23    /// Individual gradient updates.
24    pub gradients: Vec<GradientUpdate>,
25}
26
27#[derive(Deserialize)]
28struct GradientBatchWire {
29    #[serde(deserialize_with = "de_session_id")]
30    session_id: String,
31    batch_id: u64,
32    #[serde(deserialize_with = "de_gradient_rows")]
33    gradients: Vec<GradientUpdate>,
34}
35
36impl TryFrom<GradientBatchWire> for GradientBatch {
37    type Error = ValidationError;
38
39    fn try_from(wire: GradientBatchWire) -> Result<Self, Self::Error> {
40        let batch = GradientBatch {
41            session_id: wire.session_id,
42            batch_id: wire.batch_id,
43            gradients: wire.gradients,
44        };
45        batch.validate()?;
46        Ok(batch)
47    }
48}
49
50impl Validate for GradientBatch {
51    fn validate_with(&self, limits: ProtocolLimits) -> Result<(), ValidationError> {
52        check_string("session_id", &self.session_id, limits)?;
53        check_count("gradients", self.gradients.len(), limits.max_gradients)?;
54        let mut total = 0;
55        add_to_total(&mut total, self.gradients.len(), limits, "aggregate")?;
56        accumulate_gradient_rows(&self.gradients, limits, &mut total)?;
57        check_unique_strings("gradients", &self.gradients, |row| row.layer_id.as_str())?;
58        Ok(())
59    }
60}
61
62fn accumulate_gradient_rows(
63    rows: &[GradientUpdate],
64    limits: ProtocolLimits,
65    total: &mut usize,
66) -> Result<(), ValidationError> {
67    for (index, update) in rows.iter().enumerate() {
68        update
69            .validate_with(limits)
70            .map_err(|err| prefix_path(err, &format!("gradients[{index}]")))?;
71        add_to_total(total, update.gradients.len(), limits, "aggregate")?;
72        if let Some(trace) = &update.eligibility_trace {
73            add_to_total(total, trace.len(), limits, "aggregate")?;
74        }
75    }
76    Ok(())
77}
78
79/// Individual gradient update for a specific layer or parameter.
80#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
81#[serde(try_from = "GradientUpdateWire")]
82pub struct GradientUpdate {
83    /// Target layer identifier.
84    pub layer_id: String,
85    /// Gradient values (flattened or sparse representation).
86    pub gradients: Vec<f32>,
87    /// Optional eligibility trace for E-prop algorithms.
88    pub eligibility_trace: Option<Vec<f32>>,
89}
90
91#[derive(Deserialize)]
92struct GradientUpdateWire {
93    #[serde(deserialize_with = "de_layer_id")]
94    layer_id: String,
95    #[serde(deserialize_with = "de_gradient_values")]
96    gradients: Vec<f32>,
97    #[serde(default, deserialize_with = "de_eligibility_trace")]
98    eligibility_trace: Option<Vec<f32>>,
99}
100
101impl TryFrom<GradientUpdateWire> for GradientUpdate {
102    type Error = ValidationError;
103
104    fn try_from(wire: GradientUpdateWire) -> Result<Self, Self::Error> {
105        let update = GradientUpdate {
106            layer_id: wire.layer_id,
107            gradients: wire.gradients,
108            eligibility_trace: wire.eligibility_trace,
109        };
110        update.validate()?;
111        Ok(update)
112    }
113}
114
115impl Validate for GradientUpdate {
116    fn validate_with(&self, limits: ProtocolLimits) -> Result<(), ValidationError> {
117        check_string("layer_id", &self.layer_id, limits)?;
118        check_count("gradients", self.gradients.len(), limits.max_channel_values)?;
119        check_finite_slice("gradients", &self.gradients)?;
120        if let Some(trace) = &self.eligibility_trace {
121            check_count("eligibility_trace", trace.len(), limits.max_channel_values)?;
122            check_finite_slice("eligibility_trace", trace)?;
123        }
124        Ok(())
125    }
126}
corpus_ipc/models/gradients.rs

corpus_ipc/models/
gradients.rs
