ipfrs_tensorlogic/
gradient_clipper.rs

1//! Gradient clipping strategies for distributed tensor learning.
2//!
3//! Provides norm clipping and value clipping to prevent gradient explosion
4//! during distributed tensor training.
5
6/// Strategy for clipping gradients.
7#[derive(Clone, Debug, PartialEq)]
8pub enum ClippingStrategy {
9    /// Clip all gradients so that their global L2 norm is at most `max_norm`.
10    GlobalNorm {
11        /// Maximum allowed global L2 norm.
12        max_norm: f64,
13    },
14    /// Clip each gradient tensor independently so its L2 norm is at most `max_norm`.
15    PerTensorNorm {
16        /// Maximum allowed per-tensor L2 norm.
17        max_norm: f64,
18    },
19    /// Clamp every scalar value in every tensor to the range `[min, max]`.
20    ValueClip {
21        /// Minimum allowed value.
22        min: f64,
23        /// Maximum allowed value.
24        max: f64,
25    },
26    /// Running EMA of the global norm; clip when current norm > EMA * 1.5.
27    ///
28    /// The EMA is updated as: `ema = momentum * ema + (1 - momentum) * global_norm`.
29    /// On the very first call the EMA is bootstrapped to the current global norm so
30    /// no spurious clipping occurs on the initial step.
31    Adaptive {
32        /// Target norm used to scale clipping (the clip threshold is `ema * 1.5`).
33        target_norm: f64,
34        /// EMA momentum coefficient (should be in `[0, 1)`).
35        momentum: f64,
36    },
37}
38
39// ─── GradientTensor ──────────────────────────────────────────────────────────
40
41/// A single gradient tensor identified by a unique id.
42#[derive(Clone, Debug)]
43pub struct GradientTensor {
44    /// Unique identifier for this tensor.
45    pub tensor_id: u64,
46    /// The gradient values.
47    pub values: Vec<f64>,
48}
49
50impl GradientTensor {
51    /// Compute the L2 norm (Euclidean length) of the gradient values.
52    ///
53    /// Returns `0.0` for an empty tensor.
54    pub fn l2_norm(&self) -> f64 {
55        if self.values.is_empty() {
56            return 0.0;
57        }
58        let sum_sq: f64 = self.values.iter().map(|v| v * v).sum();
59        sum_sq.sqrt()
60    }
61
62    /// Return the maximum absolute value among all elements.
63    ///
64    /// Returns `0.0` for an empty tensor.
65    pub fn max_abs_value(&self) -> f64 {
66        self.values.iter().map(|v| v.abs()).fold(0.0_f64, f64::max)
67    }
68}
69
70// ─── ClippingResult ──────────────────────────────────────────────────────────
71
72/// The result of a clipping operation on a single tensor.
73#[derive(Clone, Debug)]
74pub struct ClippingResult {
75    /// The id of the tensor that was (possibly) clipped.
76    pub tensor_id: u64,
77    /// L2 norm of the tensor **before** clipping.
78    pub original_norm: f64,
79    /// L2 norm of the tensor **after** clipping.
80    pub clipped_norm: f64,
81    /// `true` if any values were actually changed by the clipper.
82    pub was_clipped: bool,
83}
84
85// ─── ClipperStats ────────────────────────────────────────────────────────────
86
87/// Cumulative statistics for a [`TensorGradientClipper`].
88#[derive(Clone, Debug, Default)]
89pub struct ClipperStats {
90    /// Number of times [`TensorGradientClipper::clip`] has been called.
91    pub total_clip_calls: u64,
92    /// Total number of tensors processed across all clip calls.
93    pub total_tensors_processed: u64,
94    /// Number of tensors for which clipping was actually applied.
95    pub total_clipped: u64,
96    /// Running mean of `clipped_norm / original_norm` for clipped tensors.
97    ///
98    /// `1.0` when no tensor has been clipped yet.
99    pub avg_clip_ratio: f64,
100}
101
102// ─── TensorGradientClipper ───────────────────────────────────────────────────
103
104/// Applies gradient-clipping strategies to collections of [`GradientTensor`]s.
105pub struct TensorGradientClipper {
106    /// The clipping strategy in use.
107    pub strategy: ClippingStrategy,
108    /// Cumulative statistics.
109    pub stats: ClipperStats,
110    /// EMA of the global norm (used only by [`ClippingStrategy::Adaptive`]).
111    pub ema_norm: f64,
112}
113
114impl TensorGradientClipper {
115    /// Create a new clipper with the given strategy and zeroed statistics.
116    pub fn new(strategy: ClippingStrategy) -> Self {
117        Self {
118            strategy,
119            stats: ClipperStats {
120                avg_clip_ratio: 1.0,
121                ..ClipperStats::default()
122            },
123            ema_norm: 0.0,
124        }
125    }
126
127    /// Apply the configured clipping strategy to `tensors` in-place.
128    ///
129    /// Returns one [`ClippingResult`] per input tensor.
130    pub fn clip(&mut self, tensors: &mut [GradientTensor]) -> Vec<ClippingResult> {
131        self.stats.total_clip_calls += 1;
132        self.stats.total_tensors_processed += tensors.len() as u64;
133
134        let results = match &self.strategy.clone() {
135            ClippingStrategy::GlobalNorm { max_norm } => self.apply_global_norm(tensors, *max_norm),
136            ClippingStrategy::PerTensorNorm { max_norm } => {
137                self.apply_per_tensor_norm(tensors, *max_norm)
138            }
139            ClippingStrategy::ValueClip { min, max } => self.apply_value_clip(tensors, *min, *max),
140            ClippingStrategy::Adaptive { momentum, .. } => {
141                let momentum = *momentum;
142                self.apply_adaptive(tensors, momentum)
143            }
144        };
145
146        // Update stats for each result
147        for result in &results {
148            if result.was_clipped {
149                self.stats.total_clipped += 1;
150                let ratio = if result.original_norm > 0.0 {
151                    result.clipped_norm / result.original_norm
152                } else {
153                    1.0
154                };
155                // Update running mean of clip ratio for clipped tensors
156                let n = self.stats.total_clipped as f64;
157                self.stats.avg_clip_ratio =
158                    self.stats.avg_clip_ratio + (ratio - self.stats.avg_clip_ratio) / n;
159            }
160        }
161
162        results
163    }
164
165    /// Reset all statistics and the EMA norm to their initial state.
166    pub fn reset_stats(&mut self) {
167        self.stats = ClipperStats {
168            avg_clip_ratio: 1.0,
169            ..ClipperStats::default()
170        };
171        self.ema_norm = 0.0;
172    }
173
174    /// Return a reference to the current statistics.
175    pub fn stats(&self) -> &ClipperStats {
176        &self.stats
177    }
178
179    // ── private helpers ──────────────────────────────────────────────────────
180
181    fn apply_global_norm(
182        &self,
183        tensors: &mut [GradientTensor],
184        max_norm: f64,
185    ) -> Vec<ClippingResult> {
186        // Compute global norm = sqrt(sum of all per-tensor squared norms)
187        let sum_sq: f64 = tensors.iter().map(|t| t.l2_norm().powi(2)).sum();
188        let global_norm = sum_sq.sqrt();
189
190        if global_norm > max_norm && global_norm > 0.0 {
191            let scale = max_norm / global_norm;
192            tensors.iter_mut().for_each(|t| {
193                t.values.iter_mut().for_each(|v| *v *= scale);
194            });
195            tensors
196                .iter()
197                .map(|t| {
198                    // After scaling: tensor_norm * scale
199                    let original = t.l2_norm() / scale; // reverse-engineer pre-clip norm
200                    let clipped = t.l2_norm();
201                    ClippingResult {
202                        tensor_id: t.tensor_id,
203                        original_norm: original,
204                        clipped_norm: clipped,
205                        was_clipped: true,
206                    }
207                })
208                .collect()
209        } else {
210            tensors
211                .iter()
212                .map(|t| {
213                    let norm = t.l2_norm();
214                    ClippingResult {
215                        tensor_id: t.tensor_id,
216                        original_norm: norm,
217                        clipped_norm: norm,
218                        was_clipped: false,
219                    }
220                })
221                .collect()
222        }
223    }
224
225    fn apply_per_tensor_norm(
226        &self,
227        tensors: &mut [GradientTensor],
228        max_norm: f64,
229    ) -> Vec<ClippingResult> {
230        tensors
231            .iter_mut()
232            .map(|t| {
233                let original_norm = t.l2_norm();
234                if original_norm > max_norm && original_norm > 0.0 {
235                    let scale = max_norm / original_norm;
236                    t.values.iter_mut().for_each(|v| *v *= scale);
237                    let clipped_norm = t.l2_norm();
238                    ClippingResult {
239                        tensor_id: t.tensor_id,
240                        original_norm,
241                        clipped_norm,
242                        was_clipped: true,
243                    }
244                } else {
245                    ClippingResult {
246                        tensor_id: t.tensor_id,
247                        original_norm,
248                        clipped_norm: original_norm,
249                        was_clipped: false,
250                    }
251                }
252            })
253            .collect()
254    }
255
256    fn apply_value_clip(
257        &self,
258        tensors: &mut [GradientTensor],
259        min: f64,
260        max: f64,
261    ) -> Vec<ClippingResult> {
262        tensors
263            .iter_mut()
264            .map(|t| {
265                let original_norm = t.l2_norm();
266                let mut any_changed = false;
267                t.values.iter_mut().for_each(|v| {
268                    let clamped = v.clamp(min, max);
269                    if clamped != *v {
270                        any_changed = true;
271                        *v = clamped;
272                    }
273                });
274                let clipped_norm = t.l2_norm();
275                ClippingResult {
276                    tensor_id: t.tensor_id,
277                    original_norm,
278                    clipped_norm,
279                    was_clipped: any_changed,
280                }
281            })
282            .collect()
283    }
284
285    fn apply_adaptive(
286        &mut self,
287        tensors: &mut [GradientTensor],
288        momentum: f64,
289    ) -> Vec<ClippingResult> {
290        const SPIKE_THRESHOLD: f64 = 1.5;
291
292        // Compute current global norm
293        let sum_sq: f64 = tensors.iter().map(|t| t.l2_norm().powi(2)).sum();
294        let global_norm = sum_sq.sqrt();
295
296        // Bootstrap EMA on first call
297        if self.ema_norm == 0.0 {
298            self.ema_norm = global_norm;
299        } else {
300            self.ema_norm = momentum * self.ema_norm + (1.0 - momentum) * global_norm;
301        }
302
303        let clip_threshold = self.ema_norm * SPIKE_THRESHOLD;
304
305        if global_norm > clip_threshold && global_norm > 0.0 {
306            // Apply global norm clip to `clip_threshold`
307            let scale = clip_threshold / global_norm;
308            tensors.iter_mut().for_each(|t| {
309                t.values.iter_mut().for_each(|v| *v *= scale);
310            });
311            tensors
312                .iter()
313                .map(|t| {
314                    let clipped_norm = t.l2_norm();
315                    let original_norm = clipped_norm / scale;
316                    ClippingResult {
317                        tensor_id: t.tensor_id,
318                        original_norm,
319                        clipped_norm,
320                        was_clipped: true,
321                    }
322                })
323                .collect()
324        } else {
325            tensors
326                .iter()
327                .map(|t| {
328                    let norm = t.l2_norm();
329                    ClippingResult {
330                        tensor_id: t.tensor_id,
331                        original_norm: norm,
332                        clipped_norm: norm,
333                        was_clipped: false,
334                    }
335                })
336                .collect()
337        }
338    }
339}
340
341// ─── Tests ────────────────────────────────────────────────────────────────────
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346
347    const EPS: f64 = 1e-9;
348
349    fn make_tensor(id: u64, values: Vec<f64>) -> GradientTensor {
350        GradientTensor {
351            tensor_id: id,
352            values,
353        }
354    }
355
356    // ── GradientTensor helpers ────────────────────────────────────────────────
357
358    #[test]
359    fn test_l2_norm_empty() {
360        let t = make_tensor(0, vec![]);
361        assert!((t.l2_norm() - 0.0).abs() < EPS);
362    }
363
364    #[test]
365    fn test_l2_norm_single() {
366        let t = make_tensor(1, vec![3.0]);
367        assert!((t.l2_norm() - 3.0).abs() < EPS);
368    }
369
370    #[test]
371    fn test_l2_norm_pythagorean() {
372        // 3-4-5 triple
373        let t = make_tensor(2, vec![3.0, 4.0]);
374        assert!((t.l2_norm() - 5.0).abs() < EPS);
375    }
376
377    #[test]
378    fn test_l2_norm_negative_values() {
379        let t = make_tensor(3, vec![-3.0, -4.0]);
380        assert!((t.l2_norm() - 5.0).abs() < EPS);
381    }
382
383    #[test]
384    fn test_max_abs_value_empty() {
385        let t = make_tensor(4, vec![]);
386        assert!((t.max_abs_value() - 0.0).abs() < EPS);
387    }
388
389    #[test]
390    fn test_max_abs_value_mixed() {
391        let t = make_tensor(5, vec![-10.0, 5.0, 3.0]);
392        assert!((t.max_abs_value() - 10.0).abs() < EPS);
393    }
394
395    #[test]
396    fn test_max_abs_value_all_negative() {
397        let t = make_tensor(6, vec![-1.0, -2.0, -0.5]);
398        assert!((t.max_abs_value() - 2.0).abs() < EPS);
399    }
400
401    // ── GlobalNorm ────────────────────────────────────────────────────────────
402
403    #[test]
404    fn test_global_norm_no_clip() {
405        let mut clipper =
406            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 10.0 });
407        let mut tensors = vec![make_tensor(1, vec![3.0, 4.0])]; // norm = 5
408        let results = clipper.clip(&mut tensors);
409        assert_eq!(results.len(), 1);
410        assert!(!results[0].was_clipped);
411        assert!((results[0].original_norm - 5.0).abs() < EPS);
412        assert!((results[0].clipped_norm - 5.0).abs() < EPS);
413    }
414
415    #[test]
416    fn test_global_norm_clip_proportionally() {
417        let mut clipper =
418            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 1.0 });
419        // global norm = sqrt(9+16) = 5; scale = 1/5
420        let mut tensors = vec![make_tensor(1, vec![3.0, 4.0])];
421        let results = clipper.clip(&mut tensors);
422        assert!(results[0].was_clipped);
423        // After clip, norm should be 1.0
424        let norm_after = tensors[0].l2_norm();
425        assert!((norm_after - 1.0).abs() < 1e-9);
426    }
427
428    #[test]
429    fn test_global_norm_clip_multi_tensor_proportional() {
430        let mut clipper =
431            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 5.0 });
432        // global norm = sqrt(9+16+25) = sqrt(50) ≈ 7.071; scale = 5/7.071
433        let mut tensors = vec![make_tensor(1, vec![3.0, 4.0]), make_tensor(2, vec![5.0])];
434        let results = clipper.clip(&mut tensors);
435        assert!(results[0].was_clipped);
436        assert!(results[1].was_clipped);
437        // Global norm after clip should equal max_norm
438        let new_global: f64 = tensors
439            .iter()
440            .map(|t| t.l2_norm().powi(2))
441            .sum::<f64>()
442            .sqrt();
443        assert!((new_global - 5.0).abs() < 1e-9);
444    }
445
446    #[test]
447    fn test_global_norm_exactly_at_threshold() {
448        let mut clipper =
449            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 5.0 });
450        let mut tensors = vec![make_tensor(1, vec![3.0, 4.0])]; // norm = 5 exactly
451        let results = clipper.clip(&mut tensors);
452        assert!(!results[0].was_clipped);
453    }
454
455    #[test]
456    fn test_global_norm_empty_tensor_list() {
457        let mut clipper =
458            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 1.0 });
459        let mut tensors: Vec<GradientTensor> = vec![];
460        let results = clipper.clip(&mut tensors);
461        assert!(results.is_empty());
462    }
463
464    // ── PerTensorNorm ─────────────────────────────────────────────────────────
465
466    #[test]
467    fn test_per_tensor_norm_clips_independently() {
468        let mut clipper =
469            TensorGradientClipper::new(ClippingStrategy::PerTensorNorm { max_norm: 3.0 });
470        let mut tensors = vec![
471            make_tensor(1, vec![3.0, 4.0]), // norm=5, will be clipped
472            make_tensor(2, vec![1.0, 2.0]), // norm≈2.24, will not be clipped
473        ];
474        let results = clipper.clip(&mut tensors);
475        assert!(results[0].was_clipped);
476        assert!(!results[1].was_clipped);
477        // Tensor 1 norm should be 3.0
478        assert!((tensors[0].l2_norm() - 3.0).abs() < 1e-9);
479        // Tensor 2 unchanged
480        assert!((tensors[1].values[0] - 1.0).abs() < EPS);
481    }
482
483    #[test]
484    fn test_per_tensor_norm_no_clip_when_under() {
485        let mut clipper =
486            TensorGradientClipper::new(ClippingStrategy::PerTensorNorm { max_norm: 10.0 });
487        let mut tensors = vec![make_tensor(1, vec![1.0, 1.0])];
488        let results = clipper.clip(&mut tensors);
489        assert!(!results[0].was_clipped);
490    }
491
492    #[test]
493    fn test_per_tensor_norm_scale_correctness() {
494        let mut clipper =
495            TensorGradientClipper::new(ClippingStrategy::PerTensorNorm { max_norm: 1.0 });
496        let mut tensors = vec![make_tensor(1, vec![0.0, 5.0])]; // norm=5
497        clipper.clip(&mut tensors);
498        // After clip, values should be [0.0, 1.0]
499        assert!((tensors[0].values[0] - 0.0).abs() < EPS);
500        assert!((tensors[0].values[1] - 1.0).abs() < 1e-9);
501    }
502
503    // ── ValueClip ─────────────────────────────────────────────────────────────
504
505    #[test]
506    fn test_value_clip_clamps_values() {
507        let mut clipper = TensorGradientClipper::new(ClippingStrategy::ValueClip {
508            min: -1.0,
509            max: 1.0,
510        });
511        let mut tensors = vec![make_tensor(1, vec![-5.0, 0.5, 3.0])];
512        let results = clipper.clip(&mut tensors);
513        assert!(results[0].was_clipped);
514        assert!((tensors[0].values[0] - (-1.0)).abs() < EPS);
515        assert!((tensors[0].values[1] - 0.5).abs() < EPS);
516        assert!((tensors[0].values[2] - 1.0).abs() < EPS);
517    }
518
519    #[test]
520    fn test_value_clip_not_clipped_when_in_range() {
521        let mut clipper = TensorGradientClipper::new(ClippingStrategy::ValueClip {
522            min: -5.0,
523            max: 5.0,
524        });
525        let mut tensors = vec![make_tensor(1, vec![-1.0, 0.0, 2.5])];
526        let results = clipper.clip(&mut tensors);
527        assert!(!results[0].was_clipped);
528    }
529
530    #[test]
531    fn test_value_clip_norm_changes() {
532        let mut clipper =
533            TensorGradientClipper::new(ClippingStrategy::ValueClip { min: 0.0, max: 1.0 });
534        let mut tensors = vec![make_tensor(1, vec![2.0, 2.0])];
535        let results = clipper.clip(&mut tensors);
536        // original norm = sqrt(8) ≈ 2.828
537        assert!((results[0].original_norm - 8_f64.sqrt()).abs() < 1e-9);
538        // clipped norm = sqrt(2) ≈ 1.414
539        assert!((results[0].clipped_norm - 2_f64.sqrt()).abs() < 1e-9);
540    }
541
542    #[test]
543    fn test_value_clip_empty_list() {
544        let mut clipper = TensorGradientClipper::new(ClippingStrategy::ValueClip {
545            min: -1.0,
546            max: 1.0,
547        });
548        let mut tensors: Vec<GradientTensor> = vec![];
549        let results = clipper.clip(&mut tensors);
550        assert!(results.is_empty());
551    }
552
553    // ── Adaptive ──────────────────────────────────────────────────────────────
554
555    #[test]
556    fn test_adaptive_no_clip_on_first_call() {
557        let mut clipper = TensorGradientClipper::new(ClippingStrategy::Adaptive {
558            target_norm: 5.0,
559            momentum: 0.9,
560        });
561        let mut tensors = vec![make_tensor(1, vec![3.0, 4.0])]; // norm=5
562        let results = clipper.clip(&mut tensors);
563        // First call: EMA bootstrapped to global_norm; threshold = global_norm*1.5 > global_norm
564        assert!(!results[0].was_clipped, "First call should never clip");
565    }
566
567    #[test]
568    fn test_adaptive_clips_on_spike() {
569        let mut clipper = TensorGradientClipper::new(ClippingStrategy::Adaptive {
570            target_norm: 5.0,
571            momentum: 0.9,
572        });
573        // First call: normal gradient, EMA ~ 1.0
574        let mut tensors1 = vec![make_tensor(1, vec![1.0])];
575        clipper.clip(&mut tensors1);
576
577        // Second call: spike at 3.0 (> 1.0 * 1.5 = 1.5)
578        let mut tensors2 = vec![make_tensor(2, vec![3.0])];
579        let results2 = clipper.clip(&mut tensors2);
580        assert!(results2[0].was_clipped, "Spike should be clipped");
581        // After clip, norm should be <= ema_norm * 1.5
582    }
583
584    #[test]
585    fn test_adaptive_ema_is_updated() {
586        let mut clipper = TensorGradientClipper::new(ClippingStrategy::Adaptive {
587            target_norm: 5.0,
588            momentum: 0.5,
589        });
590        let mut tensors = vec![make_tensor(1, vec![2.0])]; // norm=2
591        clipper.clip(&mut tensors);
592        // EMA should be bootstrapped to 2.0
593        assert!((clipper.ema_norm - 2.0).abs() < EPS);
594
595        let mut tensors2 = vec![make_tensor(2, vec![4.0])]; // norm=4
596        clipper.clip(&mut tensors2);
597        // EMA = 0.5*2 + 0.5*4 = 3.0 (uses norm from second call, which was not clipped because 4 <= 2*1.5=3 is false -- 4>3, so it IS clipped)
598        // Actually 4 > 2.0*1.5=3.0, so it clips to 3.0; new global_norm passed to EMA update is 4.0 (before clip)
599        // EMA = 0.5*2 + 0.5*4 = 3.0
600        assert!((clipper.ema_norm - 3.0).abs() < EPS);
601    }
602
603    #[test]
604    fn test_adaptive_no_clip_when_below_threshold() {
605        let mut clipper = TensorGradientClipper::new(ClippingStrategy::Adaptive {
606            target_norm: 5.0,
607            momentum: 0.9,
608        });
609        // Bootstrap EMA to 10
610        let mut tensors1 = vec![make_tensor(1, vec![10.0])];
611        clipper.clip(&mut tensors1);
612
613        // Second call: norm=5 (< 10*1.5=15), should not clip
614        let mut tensors2 = vec![make_tensor(2, vec![5.0])];
615        let results = clipper.clip(&mut tensors2);
616        assert!(!results[0].was_clipped);
617    }
618
619    // ── Stats ──────────────────────────────────────────────────────────────────
620
621    #[test]
622    fn test_stats_total_clip_calls() {
623        let mut clipper =
624            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 10.0 });
625        let mut t = vec![make_tensor(1, vec![1.0])];
626        clipper.clip(&mut t);
627        clipper.clip(&mut t);
628        assert_eq!(clipper.stats().total_clip_calls, 2);
629    }
630
631    #[test]
632    fn test_stats_total_tensors_processed() {
633        let mut clipper =
634            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 10.0 });
635        let mut tensors = vec![make_tensor(1, vec![1.0]), make_tensor(2, vec![2.0])];
636        clipper.clip(&mut tensors);
637        assert_eq!(clipper.stats().total_tensors_processed, 2);
638        clipper.clip(&mut tensors);
639        assert_eq!(clipper.stats().total_tensors_processed, 4);
640    }
641
642    #[test]
643    fn test_stats_total_clipped_counts_correctly() {
644        let mut clipper =
645            TensorGradientClipper::new(ClippingStrategy::PerTensorNorm { max_norm: 3.0 });
646        let mut tensors = vec![
647            make_tensor(1, vec![3.0, 4.0]), // norm=5, clipped
648            make_tensor(2, vec![1.0]),      // norm=1, not clipped
649        ];
650        clipper.clip(&mut tensors);
651        assert_eq!(clipper.stats().total_clipped, 1);
652    }
653
654    #[test]
655    fn test_stats_avg_clip_ratio_when_no_clipping() {
656        let mut clipper =
657            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 100.0 });
658        let mut tensors = vec![make_tensor(1, vec![1.0])];
659        clipper.clip(&mut tensors);
660        // No clipping => avg_clip_ratio stays 1.0
661        assert!((clipper.stats().avg_clip_ratio - 1.0).abs() < EPS);
662    }
663
664    #[test]
665    fn test_stats_avg_clip_ratio_running_mean() {
666        let mut clipper =
667            TensorGradientClipper::new(ClippingStrategy::PerTensorNorm { max_norm: 1.0 });
668        // First clipped tensor: original_norm=5, clipped_norm=1 => ratio=0.2
669        let mut t1 = vec![make_tensor(1, vec![0.0, 5.0])];
670        clipper.clip(&mut t1);
671        assert!((clipper.stats().avg_clip_ratio - 0.2).abs() < 1e-6);
672
673        // Second clipped tensor: original_norm=10, clipped_norm=1 => ratio=0.1
674        // running mean = (0.2 + 0.1)/2 = 0.15
675        let mut t2 = vec![make_tensor(2, vec![0.0, 10.0])];
676        clipper.clip(&mut t2);
677        assert!((clipper.stats().avg_clip_ratio - 0.15).abs() < 1e-6);
678    }
679
680    #[test]
681    fn test_reset_stats() {
682        let mut clipper =
683            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 1.0 });
684        let mut tensors = vec![make_tensor(1, vec![5.0])];
685        clipper.clip(&mut tensors);
686        clipper.reset_stats();
687        assert_eq!(clipper.stats().total_clip_calls, 0);
688        assert_eq!(clipper.stats().total_tensors_processed, 0);
689        assert_eq!(clipper.stats().total_clipped, 0);
690        assert!((clipper.stats().avg_clip_ratio - 1.0).abs() < EPS);
691        assert!((clipper.ema_norm - 0.0).abs() < EPS);
692    }
693
694    #[test]
695    fn test_empty_tensor_values_l2_norm() {
696        let t = make_tensor(99, vec![]);
697        assert!((t.l2_norm() - 0.0).abs() < EPS);
698    }
699
700    #[test]
701    fn test_global_norm_single_zero_tensor() {
702        let mut clipper =
703            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 1.0 });
704        let mut tensors = vec![make_tensor(1, vec![0.0, 0.0])];
705        let results = clipper.clip(&mut tensors);
706        // global_norm=0, no scaling
707        assert!(!results[0].was_clipped);
708    }
709
710    #[test]
711    fn test_per_tensor_norm_zero_norm_no_clip() {
712        let mut clipper =
713            TensorGradientClipper::new(ClippingStrategy::PerTensorNorm { max_norm: 1.0 });
714        let mut tensors = vec![make_tensor(1, vec![0.0])];
715        let results = clipper.clip(&mut tensors);
716        assert!(!results[0].was_clipped);
717    }
718
719    #[test]
720    fn test_value_clip_boundary_values_not_clipped() {
721        let mut clipper = TensorGradientClipper::new(ClippingStrategy::ValueClip {
722            min: -1.0,
723            max: 1.0,
724        });
725        let mut tensors = vec![make_tensor(1, vec![-1.0, 1.0])];
726        let results = clipper.clip(&mut tensors);
727        assert!(!results[0].was_clipped);
728    }
729
730    #[test]
731    fn test_adaptive_multiple_stable_calls_no_clip() {
732        let mut clipper = TensorGradientClipper::new(ClippingStrategy::Adaptive {
733            target_norm: 5.0,
734            momentum: 0.9,
735        });
736        for i in 0..5 {
737            let mut tensors = vec![make_tensor(i, vec![1.0, 1.0])]; // norm≈1.414 each time
738            let results = clipper.clip(&mut tensors);
739            assert!(
740                !results[0].was_clipped,
741                "Stable gradients should not be clipped (call {i})"
742            );
743        }
744    }
745
746    #[test]
747    fn test_clipping_result_fields() {
748        let mut clipper =
749            TensorGradientClipper::new(ClippingStrategy::GlobalNorm { max_norm: 5.0 });
750        let mut tensors = vec![make_tensor(42, vec![3.0, 4.0])]; // norm=5, no clip
751        let results = clipper.clip(&mut tensors);
752        assert_eq!(results[0].tensor_id, 42);
753    }
754}
ipfrs_tensorlogic/gradient_clipper.rs

ipfrs_tensorlogic/
gradient_clipper.rs
