ruda-tensor 0.21.74

Ruda tensor operations, backend contracts, and framework-facing primitives.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

/// Policy carried by an autodiff tensor for retaining or recomputing forward values.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GradientCheckpointingStrategy {
    /// Retain the required forward values rather than recomputing operations.
    Disabled,
    /// Recompute memory-bound operations when all required parents can be checkpointed.
    Balanced,
    /// A statically supplied backend strategy, identified by its Rust type name.
    Custom(&'static str),
}

/// Runtime autodiff association, independent of gradient retention and tape availability.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AutodiffContext {
    /// An ordinary constant with no graph association.
    Disabled,
    /// Autodiff is enabled with this checkpointing policy.
    Enabled(GradientCheckpointingStrategy),
}

impl AutodiffContext {
    /// Merge operation-local contexts without changing either operand.
    pub fn merge(self, other: Self) -> Self {
        match (self, other) {
            (Self::Disabled, context) | (context, Self::Disabled) => context,
            (Self::Enabled(left), Self::Enabled(right)) => {
                assert_eq!(left, right, "Cannot combine different gradient-checkpointing strategies");
                self
            }
        }
    }
}