tenferro_ops/
input_key.rs

1#[derive(Clone, Debug, Hash, PartialEq, Eq)]
2pub enum TensorInputKey {
3    User {
4        id: u64,
5    },
6    #[cfg(feature = "autodiff")]
7    Tangent {
8        of: Box<TensorInputKey>,
9        pass: u64,
10    },
11}
12
13impl TensorInputKey {
14    /// Returns `true` when this key names an AD tangent input.
15    ///
16    /// # Examples
17    ///
18    /// ```rust
19    /// use tenferro_ops::input_key::TensorInputKey;
20    ///
21    /// let key = TensorInputKey::User { id: 0 };
22    /// assert!(!key.is_tangent());
23    /// ```
24    pub fn is_tangent(&self) -> bool {
25        match self {
26            TensorInputKey::User { .. } => false,
27            #[cfg(feature = "autodiff")]
28            TensorInputKey::Tangent { .. } => true,
29        }
30    }
31
32    /// Returns the user input key that owns this input's concrete primal data.
33    ///
34    /// For non-AD keys this returns `self`; for tangent keys it recursively
35    /// follows the `of` chain to the original user input.
36    ///
37    /// # Examples
38    ///
39    /// ```rust
40    /// use tenferro_ops::input_key::TensorInputKey;
41    ///
42    /// let key = TensorInputKey::User { id: 0 };
43    /// assert_eq!(key.primal_root(), &key);
44    /// ```
45    pub fn primal_root(&self) -> &Self {
46        match self {
47            TensorInputKey::User { .. } => self,
48            #[cfg(feature = "autodiff")]
49            TensorInputKey::Tangent { of, .. } => of.primal_root(),
50        }
51    }
52
53    #[cfg(feature = "autodiff")]
54    pub fn tangent_of(&self, pass: u64) -> Self {
55        TensorInputKey::Tangent {
56            of: Box::new(self.clone()),
57            pass,
58        }
59    }
60}
tenferro_ops/input_key.rs

tenferro_ops/
input_key.rs
