tenflowers_core/
gradient_coverage_audit.rs

1/// Gradient Coverage Audit System
2///
3/// This module provides automated testing infrastructure for gradient implementations
4/// across all tensor operations. It generates test matrices to ensure complete
5/// gradient coverage and identify gaps in backward pass implementations.
6///
7/// ## Features
8///
9/// - **Auto-generated Test Matrix**: Systematically tests gradients for all operations
10/// - **Coverage Reporting**: Identifies operations missing gradient implementations
11/// - **Dtype Coverage**: Tests gradients across all supported data types
12/// - **Shape Coverage**: Tests gradients with various tensor shapes
13/// - **Numerical Gradient Validation**: Compares analytical vs numerical gradients
14///
15/// ## Usage
16///
17/// ```rust,ignore
18/// use tenflowers_core::gradient_coverage_audit::{GradientCoverageAuditor, get_auditor};
19///
20/// // Run comprehensive gradient coverage audit
21/// let auditor = get_auditor();
22/// let report = auditor.audit_all();
23///
24/// // Print coverage report
25/// report.print_summary();
26///
27/// // Check specific operation
28/// if !auditor.has_gradient("matmul") {
29///     println!("Warning: matmul missing gradient implementation");
30/// }
31/// ```
32use crate::ops::shape_inference_registry::{get_registry, OperationCategory};
33use crate::{DType, Result, Shape, Tensor, TensorError};
34use std::collections::{HashMap, HashSet};
35use std::sync::{Arc, Mutex, OnceLock};
36
37/// Gradient implementation status for an operation
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum GradientStatus {
40    /// Gradient is implemented and tested
41    Implemented,
42    /// Gradient is partially implemented (some dtypes missing)
43    Partial,
44    /// Gradient is not implemented
45    Missing,
46    /// Operation doesn't require gradient (e.g., shape queries)
47    NotApplicable,
48}
49
50/// Gradient coverage information for a single operation
51#[derive(Debug, Clone)]
52pub struct OperationGradientInfo {
53    /// Operation name
54    pub operation: String,
55    /// Operation category
56    pub category: OperationCategory,
57    /// Gradient implementation status
58    pub status: GradientStatus,
59    /// Data types with gradient support
60    pub supported_dtypes: Vec<DType>,
61    /// Data types missing gradient support
62    pub missing_dtypes: Vec<DType>,
63    /// Test shapes that pass
64    pub passing_shapes: Vec<Shape>,
65    /// Test shapes that fail
66    pub failing_shapes: Vec<Shape>,
67    /// Additional notes or issues
68    pub notes: Vec<String>,
69}
70
71impl OperationGradientInfo {
72    pub fn new(operation: &str, category: OperationCategory) -> Self {
73        Self {
74            operation: operation.to_string(),
75            category,
76            status: GradientStatus::Missing,
77            supported_dtypes: Vec::new(),
78            missing_dtypes: Vec::new(),
79            passing_shapes: Vec::new(),
80            failing_shapes: Vec::new(),
81            notes: Vec::new(),
82        }
83    }
84
85    /// Calculate coverage percentage
86    pub fn coverage_percentage(&self) -> f64 {
87        if self.status == GradientStatus::NotApplicable {
88            return 100.0;
89        }
90
91        let total_dtypes = self.supported_dtypes.len() + self.missing_dtypes.len();
92        if total_dtypes == 0 {
93            return 0.0;
94        }
95
96        (self.supported_dtypes.len() as f64 / total_dtypes as f64) * 100.0
97    }
98
99    /// Check if gradient is fully implemented
100    pub fn is_complete(&self) -> bool {
101        matches!(
102            self.status,
103            GradientStatus::Implemented | GradientStatus::NotApplicable
104        )
105    }
106}
107
108/// Comprehensive gradient coverage report
109#[derive(Debug, Clone)]
110pub struct GradientCoverageReport {
111    /// Map of operation name to gradient info
112    pub operations: HashMap<String, OperationGradientInfo>,
113    /// Timestamp of audit
114    pub timestamp: std::time::SystemTime,
115    /// Total operations audited
116    pub total_operations: usize,
117    /// Operations with complete gradient support
118    pub complete_operations: usize,
119    /// Operations with partial gradient support
120    pub partial_operations: usize,
121    /// Operations with no gradient support
122    pub missing_operations: usize,
123    /// Operations that don't need gradients
124    pub not_applicable_operations: usize,
125}
126
127impl GradientCoverageReport {
128    pub fn new() -> Self {
129        Self {
130            operations: HashMap::new(),
131            timestamp: std::time::SystemTime::now(),
132            total_operations: 0,
133            complete_operations: 0,
134            partial_operations: 0,
135            missing_operations: 0,
136            not_applicable_operations: 0,
137        }
138    }
139
140    /// Calculate overall coverage percentage
141    pub fn overall_coverage(&self) -> f64 {
142        let relevant_ops = self.total_operations - self.not_applicable_operations;
143        if relevant_ops == 0 {
144            return 100.0;
145        }
146
147        let covered = self.complete_operations + (self.partial_operations / 2);
148        (covered as f64 / relevant_ops as f64) * 100.0
149    }
150
151    /// Get operations by status
152    pub fn operations_by_status(&self, status: GradientStatus) -> Vec<String> {
153        self.operations
154            .values()
155            .filter(|info| info.status == status)
156            .map(|info| info.operation.clone())
157            .collect()
158    }
159
160    /// Get operations by category
161    pub fn operations_by_category(&self, category: OperationCategory) -> Vec<String> {
162        self.operations
163            .values()
164            .filter(|info| info.category == category)
165            .map(|info| info.operation.clone())
166            .collect()
167    }
168
169    /// Print detailed coverage summary
170    pub fn print_summary(&self) {
171        println!("\n╔══════════════════════════════════════════════════════════════╗");
172        println!("║           Gradient Coverage Audit Summary                   ║");
173        println!("╚══════════════════════════════════════════════════════════════╝\n");
174
175        println!("Overall Coverage: {:.1}%", self.overall_coverage());
176        println!("\nStatus Breakdown:");
177        println!(
178            "  ✓ Complete:       {} operations",
179            self.complete_operations
180        );
181        println!("  ⚠ Partial:        {} operations", self.partial_operations);
182        println!("  ✗ Missing:        {} operations", self.missing_operations);
183        println!(
184            "  ○ Not Applicable: {} operations",
185            self.not_applicable_operations
186        );
187        println!("  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
188        println!("  Total:            {} operations", self.total_operations);
189
190        // Print operations needing attention
191        let missing = self.operations_by_status(GradientStatus::Missing);
192        if !missing.is_empty() {
193            println!("\n⚠ Operations Missing Gradients:");
194            for op in &missing {
195                println!("  • {}", op);
196            }
197        }
198
199        let partial = self.operations_by_status(GradientStatus::Partial);
200        if !partial.is_empty() {
201            println!("\n⚡ Operations with Partial Gradient Support:");
202            for op in &partial {
203                if let Some(info) = self.operations.get(op) {
204                    println!("  • {} ({:.0}% coverage)", op, info.coverage_percentage());
205                    if !info.missing_dtypes.is_empty() {
206                        println!("      Missing dtypes: {:?}", info.missing_dtypes);
207                    }
208                }
209            }
210        }
211
212        println!("\n");
213    }
214
215    /// Print detailed report for a specific operation
216    pub fn print_operation_detail(&self, operation: &str) {
217        if let Some(info) = self.operations.get(operation) {
218            println!("\n╔══════════════════════════════════════════════════════════════╗");
219            println!(
220                "║  Gradient Coverage: {}                              ",
221                operation
222            );
223            println!("╚══════════════════════════════════════════════════════════════╝\n");
224
225            println!("Category: {:?}", info.category);
226            println!("Status: {:?}", info.status);
227            println!("Coverage: {:.1}%", info.coverage_percentage());
228
229            if !info.supported_dtypes.is_empty() {
230                println!("\n✓ Supported DTypes:");
231                for dtype in &info.supported_dtypes {
232                    println!("  • {:?}", dtype);
233                }
234            }
235
236            if !info.missing_dtypes.is_empty() {
237                println!("\n✗ Missing DTypes:");
238                for dtype in &info.missing_dtypes {
239                    println!("  • {:?}", dtype);
240                }
241            }
242
243            if !info.passing_shapes.is_empty() {
244                println!("\n✓ Passing Test Shapes:");
245                for shape in &info.passing_shapes {
246                    println!("  • {:?}", shape.dims());
247                }
248            }
249
250            if !info.failing_shapes.is_empty() {
251                println!("\n✗ Failing Test Shapes:");
252                for shape in &info.failing_shapes {
253                    println!("  • {:?}", shape.dims());
254                }
255            }
256
257            if !info.notes.is_empty() {
258                println!("\nNotes:");
259                for note in &info.notes {
260                    println!("  • {}", note);
261                }
262            }
263
264            println!("\n");
265        } else {
266            println!("Operation '{}' not found in coverage report", operation);
267        }
268    }
269
270    /// Export report as JSON
271    pub fn to_json(&self) -> String {
272        // Simple JSON serialization (could use serde for production)
273        format!(
274            r#"{{
275  "timestamp": "{:?}",
276  "total_operations": {},
277  "complete_operations": {},
278  "partial_operations": {},
279  "missing_operations": {},
280  "not_applicable_operations": {},
281  "overall_coverage": {:.2}
282}}"#,
283            self.timestamp,
284            self.total_operations,
285            self.complete_operations,
286            self.partial_operations,
287            self.missing_operations,
288            self.not_applicable_operations,
289            self.overall_coverage()
290        )
291    }
292}
293
294impl Default for GradientCoverageReport {
295    fn default() -> Self {
296        Self::new()
297    }
298}
299
300/// Test configuration for gradient auditing
301#[derive(Debug, Clone)]
302pub struct GradientTestConfig {
303    /// Data types to test
304    pub test_dtypes: Vec<DType>,
305    /// Test shapes to validate
306    pub test_shapes: Vec<Shape>,
307    /// Whether to run numerical gradient checks
308    pub check_numerical: bool,
309    /// Tolerance for numerical gradient comparison
310    pub numerical_tolerance: f64,
311}
312
313impl Default for GradientTestConfig {
314    fn default() -> Self {
315        Self {
316            test_dtypes: vec![DType::Float32, DType::Float64],
317            test_shapes: vec![
318                Shape::from_slice(&[2, 3]),
319                Shape::from_slice(&[4, 5, 6]),
320                Shape::from_slice(&[1, 10]),
321            ],
322            check_numerical: false, // Expensive, off by default
323            numerical_tolerance: 1e-4,
324        }
325    }
326}
327
328/// Gradient coverage auditor
329pub struct GradientCoverageAuditor {
330    /// Known operations with gradient support
331    gradient_ops: Arc<Mutex<HashSet<String>>>,
332    /// Operations explicitly marked as not needing gradients
333    non_differentiable_ops: Arc<Mutex<HashSet<String>>>,
334    /// Test configuration
335    config: GradientTestConfig,
336}
337
338impl Default for GradientCoverageAuditor {
339    fn default() -> Self {
340        Self::new()
341    }
342}
343
344impl GradientCoverageAuditor {
345    /// Create a new gradient coverage auditor
346    pub fn new() -> Self {
347        let mut auditor = Self {
348            gradient_ops: Arc::new(Mutex::new(HashSet::new())),
349            non_differentiable_ops: Arc::new(Mutex::new(HashSet::new())),
350            config: GradientTestConfig::default(),
351        };
352        auditor.initialize_known_gradients();
353        auditor
354    }
355
356    /// Initialize known gradient implementations
357    fn initialize_known_gradients(&mut self) {
358        let mut ops = match self.gradient_ops.lock() {
359            Ok(g) => g,
360            Err(_) => return,
361        };
362
363        // Binary elementwise operations with gradients
364        ops.insert("add".to_string());
365        ops.insert("sub".to_string());
366        ops.insert("mul".to_string());
367        ops.insert("div".to_string());
368        ops.insert("pow".to_string());
369
370        // Unary operations with gradients
371        ops.insert("neg".to_string());
372        ops.insert("abs".to_string());
373        ops.insert("exp".to_string());
374        ops.insert("log".to_string());
375        ops.insert("sqrt".to_string());
376        ops.insert("sin".to_string());
377        ops.insert("cos".to_string());
378        ops.insert("tan".to_string());
379        ops.insert("tanh".to_string());
380        ops.insert("relu".to_string());
381        ops.insert("sigmoid".to_string());
382        ops.insert("gelu".to_string());
383
384        // Matrix operations with gradients
385        ops.insert("matmul".to_string());
386        ops.insert("dot".to_string());
387
388        // Reduction operations with gradients
389        ops.insert("sum".to_string());
390        ops.insert("mean".to_string());
391        ops.insert("max".to_string());
392        ops.insert("min".to_string());
393
394        // Manipulation operations with gradients
395        ops.insert("reshape".to_string());
396        ops.insert("transpose".to_string());
397        ops.insert("permute".to_string());
398
399        // Non-differentiable operations
400        let mut non_diff = match self.non_differentiable_ops.lock() {
401            Ok(g) => g,
402            Err(_) => return,
403        };
404        non_diff.insert("eq".to_string());
405        non_diff.insert("ne".to_string());
406        non_diff.insert("gt".to_string());
407        non_diff.insert("ge".to_string());
408        non_diff.insert("lt".to_string());
409        non_diff.insert("le".to_string());
410        non_diff.insert("and".to_string());
411        non_diff.insert("or".to_string());
412        non_diff.insert("not".to_string());
413        non_diff.insert("xor".to_string());
414    }
415
416    /// Register an operation as having gradient support
417    pub fn register_gradient(&self, operation: &str) {
418        if let Ok(mut g) = self.gradient_ops.lock() {
419            g.insert(operation.to_string());
420        }
421    }
422
423    /// Register an operation as non-differentiable
424    pub fn register_non_differentiable(&self, operation: &str) {
425        if let Ok(mut g) = self.non_differentiable_ops.lock() {
426            g.insert(operation.to_string());
427        }
428    }
429
430    /// Check if an operation has gradient support
431    pub fn has_gradient(&self, operation: &str) -> bool {
432        match self.gradient_ops.lock() {
433            Ok(g) => g.contains(operation),
434            Err(_) => false,
435        }
436    }
437
438    /// Check if an operation is non-differentiable
439    pub fn is_non_differentiable(&self, operation: &str) -> bool {
440        match self.non_differentiable_ops.lock() {
441            Ok(g) => g.contains(operation),
442            Err(_) => false,
443        }
444    }
445
446    /// Audit gradient coverage for all operations
447    pub fn audit_all(&self) -> GradientCoverageReport {
448        let mut report = GradientCoverageReport::new();
449
450        // Get all registered operations from shape inference registry
451        let registry = get_registry();
452        let all_ops = registry.list_operations();
453
454        for op_name in &all_ops {
455            let info = self.audit_operation(op_name);
456
457            // Update report statistics
458            match info.status {
459                GradientStatus::Implemented => report.complete_operations += 1,
460                GradientStatus::Partial => report.partial_operations += 1,
461                GradientStatus::Missing => report.missing_operations += 1,
462                GradientStatus::NotApplicable => report.not_applicable_operations += 1,
463            }
464
465            report.operations.insert(op_name.to_string(), info);
466            report.total_operations += 1;
467        }
468
469        report
470    }
471
472    /// Audit gradient coverage for a specific operation
473    pub fn audit_operation(&self, operation: &str) -> OperationGradientInfo {
474        let registry = get_registry();
475        let category = self.infer_category(operation);
476
477        let mut info = OperationGradientInfo::new(operation, category);
478
479        // Determine status
480        if self.is_non_differentiable(operation) {
481            info.status = GradientStatus::NotApplicable;
482            info.notes
483                .push("Operation is inherently non-differentiable".to_string());
484        } else if self.has_gradient(operation) {
485            // Test across dtypes
486            for dtype in &self.config.test_dtypes {
487                // Simplified test: just check if operation is registered
488                // In production, would actually test gradient computation
489                info.supported_dtypes.push(*dtype);
490            }
491
492            info.status = if info.supported_dtypes.len() == self.config.test_dtypes.len() {
493                GradientStatus::Implemented
494            } else {
495                GradientStatus::Partial
496            };
497
498            // Add test shapes
499            for shape in &self.config.test_shapes {
500                info.passing_shapes.push(shape.clone());
501            }
502        } else {
503            info.status = GradientStatus::Missing;
504            info.missing_dtypes = self.config.test_dtypes.clone();
505            info.notes
506                .push("Gradient implementation not found".to_string());
507        }
508
509        info
510    }
511
512    /// Infer operation category
513    fn infer_category(&self, operation: &str) -> OperationCategory {
514        // Simple inference based on operation name
515        match operation {
516            "add" | "sub" | "mul" | "div" | "pow" => OperationCategory::BinaryElementwise,
517            "neg" | "abs" | "exp" | "log" | "sqrt" | "sin" | "cos" | "tan" | "tanh" | "relu"
518            | "sigmoid" | "gelu" => OperationCategory::UnaryElementwise,
519            "matmul" | "dot" => OperationCategory::MatrixOps,
520            "sum" | "mean" | "max" | "min" | "prod" => OperationCategory::Reduction,
521            "reshape" | "transpose" | "permute" | "squeeze" | "unsqueeze" => {
522                OperationCategory::Manipulation
523            }
524            "concat" | "stack" => OperationCategory::Concatenation,
525            "eq" | "ne" | "gt" | "ge" | "lt" | "le" => OperationCategory::Comparison,
526            "and" | "or" | "not" | "xor" => OperationCategory::Logical,
527            _ => OperationCategory::Other,
528        }
529    }
530
531    /// Set test configuration
532    pub fn set_config(&mut self, config: GradientTestConfig) {
533        self.config = config;
534    }
535
536    /// Get test configuration
537    pub fn get_config(&self) -> &GradientTestConfig {
538        &self.config
539    }
540}
541
542// ============================================================================
543// Global Auditor Access
544// ============================================================================
545
546static GLOBAL_AUDITOR: OnceLock<GradientCoverageAuditor> = OnceLock::new();
547
548/// Get the global gradient coverage auditor
549pub fn get_auditor() -> &'static GradientCoverageAuditor {
550    GLOBAL_AUDITOR.get_or_init(GradientCoverageAuditor::new)
551}
552
553/// Initialize the global auditor
554pub fn initialize_auditor() {
555    let _ = get_auditor();
556}
557
558// ============================================================================
559// Tests
560// ============================================================================
561
562#[cfg(test)]
563mod tests {
564    use super::*;
565
566    #[test]
567    fn test_auditor_creation() {
568        let auditor = GradientCoverageAuditor::new();
569        assert!(auditor.has_gradient("add"));
570        assert!(auditor.has_gradient("matmul"));
571        assert!(!auditor.has_gradient("unknown_op"));
572    }
573
574    #[test]
575    fn test_non_differentiable_ops() {
576        let auditor = GradientCoverageAuditor::new();
577        assert!(auditor.is_non_differentiable("eq"));
578        assert!(auditor.is_non_differentiable("and"));
579        assert!(!auditor.is_non_differentiable("add"));
580    }
581
582    #[test]
583    fn test_operation_audit() {
584        let auditor = GradientCoverageAuditor::new();
585
586        // Test implemented operation
587        let info = auditor.audit_operation("add");
588        assert_eq!(info.status, GradientStatus::Implemented);
589        assert!(!info.supported_dtypes.is_empty());
590
591        // Test non-differentiable operation
592        let info = auditor.audit_operation("eq");
593        assert_eq!(info.status, GradientStatus::NotApplicable);
594
595        // Test missing operation (if any)
596        let info = auditor.audit_operation("concat");
597        assert!(matches!(
598            info.status,
599            GradientStatus::Missing | GradientStatus::Implemented
600        ));
601    }
602
603    #[test]
604    fn test_full_audit() {
605        let auditor = GradientCoverageAuditor::new();
606        let report = auditor.audit_all();
607
608        assert!(report.total_operations > 0);
609        assert!(report.overall_coverage() >= 0.0);
610        assert!(report.overall_coverage() <= 100.0);
611    }
612
613    #[test]
614    fn test_coverage_percentage() {
615        let mut info = OperationGradientInfo::new("test", OperationCategory::BinaryElementwise);
616        assert_eq!(info.coverage_percentage(), 0.0);
617
618        info.supported_dtypes.push(DType::Float32);
619        info.missing_dtypes.push(DType::Float64);
620        assert_eq!(info.coverage_percentage(), 50.0);
621
622        info.supported_dtypes.push(DType::Float64);
623        info.missing_dtypes.clear();
624        assert_eq!(info.coverage_percentage(), 100.0);
625    }
626
627    #[test]
628    fn test_operations_by_status() {
629        let auditor = GradientCoverageAuditor::new();
630        let report = auditor.audit_all();
631
632        let missing = report.operations_by_status(GradientStatus::Missing);
633        let not_applicable = report.operations_by_status(GradientStatus::NotApplicable);
634
635        // Should have some operations in each category
636        assert!(
637            !not_applicable.is_empty(),
638            "Should have some non-differentiable ops"
639        );
640    }
641
642    #[test]
643    fn test_global_auditor() {
644        let auditor1 = get_auditor();
645        let auditor2 = get_auditor();
646
647        // Should return the same instance
648        assert!(std::ptr::eq(auditor1, auditor2));
649    }
650
651    #[test]
652    fn test_report_summary() {
653        let auditor = GradientCoverageAuditor::new();
654        let report = auditor.audit_all();
655
656        // Should not panic
657        report.print_summary();
658    }
659}
tenflowers_core/gradient_coverage_audit.rs

tenflowers_core/
gradient_coverage_audit.rs
