tenferro_einsum/
tensordot.rs

1use std::collections::HashSet;
2
3use tenferro_runtime::error::{Error, ErrorPhase, Result};
4use tenferro_runtime::{DotGeneralConfig, TracedTensor};
5use tenferro_tensor::{ShapeMismatch, ValidationError};
6
7/// Axis specification for [`TracedTensorEinsumExt::tensordot`](crate::TracedTensorEinsumExt::tensordot)
8/// contraction sugar.
9///
10/// `Count(n)` contracts the last `n` axes of the left operand with the first
11/// `n` axes of the right operand. `Axes` contracts the explicitly paired axes
12/// in the order they are provided, accepting negative indices relative to each
13/// operand rank.
14///
15/// # Examples
16///
17/// ```
18/// use tenferro_einsum::TensorDotAxes;
19///
20/// let count = TensorDotAxes::Count(2);
21/// let explicit = TensorDotAxes::Axes {
22///     lhs: &[-1],
23///     rhs: &[0],
24/// };
25///
26/// assert_eq!(count, TensorDotAxes::Count(2));
27/// assert_ne!(count, explicit);
28/// ```
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub enum TensorDotAxes<'a> {
31    /// Contract the last `n` left axes with the first `n` right axes.
32    Count(usize),
33    /// Contract explicit left/right axis pairs.
34    Axes {
35        /// Left operand axes. Negative axes are normalized against the left rank.
36        lhs: &'a [isize],
37        /// Right operand axes. Negative axes are normalized against the right rank.
38        rhs: &'a [isize],
39    },
40}
41
42pub(crate) fn dot_general_config(
43    axes: TensorDotAxes<'_>,
44    lhs_rank: usize,
45    rhs_rank: usize,
46) -> Result<DotGeneralConfig> {
47    let (lhs_contracting_dims, rhs_contracting_dims) = match axes {
48        TensorDotAxes::Count(count) => {
49            if count > lhs_rank || count > rhs_rank {
50                return Err(Error::invalid_argument(
51                    "tensordot",
52                    ErrorPhase::GraphBuild,
53                    "axes",
54                    format!(
55                        "TensorDotAxes::Count({count}) cannot contract {count} axes \
56                         for lhs rank {lhs_rank} and rhs rank {rhs_rank}"
57                    ),
58                ));
59            }
60            ((lhs_rank - count..lhs_rank).collect(), (0..count).collect())
61        }
62        TensorDotAxes::Axes { lhs, rhs } => {
63            if lhs.len() != rhs.len() {
64                return Err(Error::invalid_argument(
65                    "tensordot",
66                    ErrorPhase::GraphBuild,
67                    "axes",
68                    format!(
69                        "tensordot explicit axes must have matching lengths, got lhs {} and rhs {}",
70                        lhs.len(),
71                        rhs.len()
72                    ),
73                ));
74            }
75            (
76                normalize_axes(lhs, lhs_rank, "lhs")?,
77                normalize_axes(rhs, rhs_rank, "rhs")?,
78            )
79        }
80    };
81
82    let config = DotGeneralConfig {
83        lhs_contracting_dims,
84        rhs_contracting_dims,
85        lhs_batch_dims: Vec::new(),
86        rhs_batch_dims: Vec::new(),
87    };
88    config
89        .validate_dims_with_ranks(lhs_rank, rhs_rank)
90        .map_err(|error| graph_build_tensor_error("tensordot", error))?;
91    Ok(config)
92}
93
94pub(crate) fn validate_concrete_contract_dims(
95    lhs_shape: &[usize],
96    rhs_shape: &[usize],
97    config: &DotGeneralConfig,
98) -> Result<()> {
99    config
100        .validate_dims_with_ranks(lhs_shape.len(), rhs_shape.len())
101        .map_err(|error| graph_build_tensor_error("tensordot", error))?;
102    for (&lhs_axis, &rhs_axis) in config
103        .lhs_contracting_dims
104        .iter()
105        .zip(config.rhs_contracting_dims.iter())
106    {
107        let lhs_dim = lhs_shape[lhs_axis];
108        let rhs_dim = rhs_shape[rhs_axis];
109        if lhs_dim != rhs_dim {
110            return Err(contracted_dims_error(lhs_axis, lhs_dim, rhs_axis, rhs_dim));
111        }
112    }
113    Ok(())
114}
115
116pub(crate) fn validate_traced_contract_dims(
117    lhs: &TracedTensor,
118    rhs: &TracedTensor,
119    config: &DotGeneralConfig,
120) -> Result<()> {
121    config
122        .validate_dims_with_ranks(lhs.rank, rhs.rank)
123        .map_err(|error| graph_build_tensor_error("tensordot", error))?;
124    for (&lhs_axis, &rhs_axis) in config
125        .lhs_contracting_dims
126        .iter()
127        .zip(config.rhs_contracting_dims.iter())
128    {
129        let lhs_dim = lhs.axis_sym_dim(lhs_axis)?;
130        let rhs_dim = rhs.axis_sym_dim(rhs_axis)?;
131        if lhs_dim == rhs_dim {
132            continue;
133        }
134        if let (Some(lhs_value), Some(rhs_value)) =
135            (lhs_dim.constant_value(), rhs_dim.constant_value())
136        {
137            if lhs_value != rhs_value {
138                return Err(contracted_dims_error(
139                    lhs_axis, lhs_value, rhs_axis, rhs_value,
140                ));
141            }
142        }
143    }
144    Ok(())
145}
146
147fn normalize_axes(axes: &[isize], rank: usize, operand: &'static str) -> Result<Vec<usize>> {
148    let mut normalized = Vec::with_capacity(axes.len());
149    let mut seen = HashSet::with_capacity(axes.len());
150    for &axis in axes {
151        let normalized_axis = normalize_axis(axis, rank, operand)?;
152        if !seen.insert(normalized_axis) {
153            return Err(Error::validation(
154                "tensordot",
155                ErrorPhase::GraphBuild,
156                ValidationError::DuplicateAxis {
157                    axis: normalized_axis,
158                    role: operand,
159                },
160            ));
161        }
162        normalized.push(normalized_axis);
163    }
164    Ok(normalized)
165}
166
167fn normalize_axis(axis: isize, rank: usize, _operand: &'static str) -> Result<usize> {
168    let rank_isize = isize::try_from(rank).map_err(|_| {
169        Error::validation(
170            "tensordot",
171            ErrorPhase::GraphBuild,
172            ValidationError::IntegerOverflow,
173        )
174    })?;
175    let normalized = if axis < 0 { rank_isize + axis } else { axis };
176    if normalized < 0 || normalized >= rank_isize {
177        return Err(Error::validation(
178            "tensordot",
179            ErrorPhase::GraphBuild,
180            ValidationError::AxisOutOfBounds {
181                axis: axis.unsigned_abs(),
182                rank,
183            },
184        ));
185    }
186    usize::try_from(normalized).map_err(|_| {
187        Error::validation(
188            "tensordot",
189            ErrorPhase::GraphBuild,
190            ValidationError::IntegerOverflow,
191        )
192    })
193}
194
195fn contracted_dims_error(
196    lhs_axis: usize,
197    lhs_dim: usize,
198    rhs_axis: usize,
199    rhs_dim: usize,
200) -> Error {
201    Error::validation(
202        "tensordot",
203        ErrorPhase::GraphBuild,
204        ShapeMismatch::ContractedDimensions {
205            lhs_axis,
206            lhs_size: lhs_dim,
207            rhs_axis,
208            rhs_size: rhs_dim,
209        }
210        .into(),
211    )
212}
213
214fn graph_build_tensor_error(op: &'static str, error: tenferro_tensor::Error) -> Error {
215    match error {
216        tenferro_tensor::Error::Validation { source, .. } => {
217            Error::validation(op, ErrorPhase::GraphBuild, source)
218        }
219        error => Error::TensorRuntime(error),
220    }
221}
222
223#[cfg(test)]
224mod tests;
tenferro_einsum/tensordot.rs

tenferro_einsum/
tensordot.rs
