tenferro_einsum/
concrete.rs

1//! Public concrete tensor einsum extension API.
2
3use tenferro_tensor::{
4    DType, DotGeneralAccumulation, Tensor, TensorBackend, TensorRead, TensorScalar, TensorWrite,
5    TypedTensor, TypedTensorView, TypedTensorWrite,
6};
7
8use crate::eager::{
9    eager_einsum_exec, eager_einsum_exec_read, eager_einsum_exec_read_into,
10    eager_einsum_exec_read_into_accum, eager_einsum_read_subscripts, eager_einsum_subscripts,
11    plan_subscripts,
12};
13use crate::TensorDotAxes;
14use crate::{ContractionTree, EinsumSubscripts, Error, Result, Subscripts};
15
16const TENSOR_EINSUM_OP: &str = "TensorEinsumExt::einsum";
17const TENSOR_EINSUM_INTO_OP: &str = "TensorEinsumIntoExt::einsum_into";
18const TENSOR_READ_EINSUM_OP: &str = "TensorReadEinsumExt::einsum_read";
19const TENSOR_READ_EINSUM_INTO_OP: &str = "TensorReadEinsumIntoExt::einsum_read_into";
20const TYPED_TENSOR_EINSUM_OP: &str = "TypedTensorEinsumExt::einsum";
21const TYPED_TENSOR_EINSUM_INTO_OP: &str = "TypedTensorEinsumIntoExt::einsum_into";
22const TYPED_TENSOR_READ_EINSUM_OP: &str = "TypedTensorReadEinsumExt::einsum_read";
23const TYPED_TENSOR_READ_EINSUM_INTO_OP: &str = "TypedTensorReadEinsumIntoExt::einsum_read_into";
24const PLAN_PREPARE_OP: &str = "ConcreteEinsumPlan::prepare";
25const PLAN_EXECUTE_OP: &str = "ConcreteEinsumPlan::execute";
26const TYPED_TENSOR_TENSORDOT_OP: &str = "TypedTensorTensordotExt::tensordot";
27
28/// Backend-explicit tensordot sugar for dtype-erased concrete tensors.
29pub trait TensorTensordotExt {
30    /// Contract this tensor with `rhs` over explicit axes or an axis count.
31    ///
32    /// # Errors
33    ///
34    /// Returns [`Error::Validation`] with `InvalidArgument`, `RankMismatch`,
35    /// `AxisOutOfBounds`, or `ShapeMismatch` for invalid axes or incompatible
36    /// contracting dimensions, or [`Error::Tensor`] when backend execution
37    /// fails.
38    fn tensordot<B: TensorBackend>(
39        &self,
40        rhs: &Tensor,
41        axes: TensorDotAxes<'_>,
42        backend: &mut B,
43    ) -> Result<Tensor>;
44}
45
46impl TensorTensordotExt for Tensor {
47    fn tensordot<B: TensorBackend>(
48        &self,
49        rhs: &Tensor,
50        axes: TensorDotAxes<'_>,
51        backend: &mut B,
52    ) -> Result<Tensor> {
53        let config =
54            crate::tensordot::dot_general_config(axes, self.shape().len(), rhs.shape().len())?;
55        crate::tensordot::validate_concrete_contract_dims(self.shape(), rhs.shape(), &config)?;
56        backend.dot_general(self, rhs, &config).map_err(Error::from)
57    }
58}
59
60/// Backend-explicit tensordot sugar for typed concrete tensors.
61pub trait TypedTensorTensordotExt<T: TensorScalar> {
62    /// Contract this tensor with `rhs` while preserving its scalar type.
63    ///
64    /// # Errors
65    ///
66    /// Returns [`Error::Validation`] with `InvalidArgument`, `RankMismatch`,
67    /// `AxisOutOfBounds`, or `ShapeMismatch` for invalid axes or incompatible
68    /// contracting dimensions, or [`Error::Tensor`] when backend execution
69    /// fails.
70    fn tensordot<B: TensorBackend>(
71        &self,
72        rhs: &TypedTensor<T>,
73        axes: TensorDotAxes<'_>,
74        backend: &mut B,
75    ) -> Result<TypedTensor<T>>;
76}
77
78impl<T: TensorScalar> TypedTensorTensordotExt<T> for TypedTensor<T> {
79    fn tensordot<B: TensorBackend>(
80        &self,
81        rhs: &TypedTensor<T>,
82        axes: TensorDotAxes<'_>,
83        backend: &mut B,
84    ) -> Result<TypedTensor<T>> {
85        let config =
86            crate::tensordot::dot_general_config(axes, self.shape().len(), rhs.shape().len())?;
87        crate::tensordot::validate_concrete_contract_dims(self.shape(), rhs.shape(), &config)?;
88        let result = backend
89            .dot_general_read(T::tensor_read(self), T::tensor_read(rhs), &config)
90            .map_err(Error::from)?;
91        into_typed_result(result, TYPED_TENSOR_TENSORDOT_OP)
92    }
93}
94
95/// Backend-explicit einsum methods for dtype-erased concrete tensors.
96///
97/// Implementations are provided for slices and fixed-size arrays of
98/// [`Tensor`] references, so both `inputs.as_slice().einsum(...)` and
99/// `[&lhs, &rhs].einsum(...)` work.
100///
101/// # Examples
102///
103/// ```
104/// use tenferro_cpu::CpuBackend;
105/// use tenferro_einsum::TensorEinsumExt;
106/// use tenferro_tensor::Tensor;
107///
108/// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
109/// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
110/// let mut backend = CpuBackend::new();
111///
112/// let out = [&lhs, &rhs].einsum("ij,jk->ik", &mut backend)?;
113/// assert_eq!(out.shape(), &[2, 4]);
114/// # Ok::<(), tenferro_einsum::Error>(())
115/// ```
116pub trait TensorEinsumExt {
117    /// Execute an einsum from string notation.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
122    /// [`Error::Validation`] with shape, rank, or dtype payloads for an invalid
123    /// contraction, or [`Error::Tensor`] for a typed backend failure.
124    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor>;
125
126    /// Execute an einsum from parsed integer-label subscripts.
127    ///
128    /// # Errors
129    ///
130    /// Returns [`Error::Validation`] with shape, rank, or dtype payloads for an
131    /// invalid contraction, or [`Error::Tensor`] for a typed backend failure.
132    fn einsum_subscripts<B: TensorBackend>(
133        &self,
134        subscripts: &EinsumSubscripts,
135        backend: &mut B,
136    ) -> Result<Tensor>;
137}
138
139impl TensorEinsumExt for [&Tensor] {
140    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
141        let subscripts = parse_subscripts(subscripts, TENSOR_EINSUM_OP)?;
142        eager_einsum_subscripts(backend, self, &subscripts).map_err(Error::from)
143    }
144
145    fn einsum_subscripts<B: TensorBackend>(
146        &self,
147        subscripts: &EinsumSubscripts,
148        backend: &mut B,
149    ) -> Result<Tensor> {
150        let subscripts = Subscripts::from(subscripts);
151        eager_einsum_subscripts(backend, self, &subscripts).map_err(Error::from)
152    }
153}
154
155impl<const N: usize> TensorEinsumExt for [&Tensor; N] {
156    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
157        self.as_slice().einsum(subscripts, backend)
158    }
159
160    fn einsum_subscripts<B: TensorBackend>(
161        &self,
162        subscripts: &EinsumSubscripts,
163        backend: &mut B,
164    ) -> Result<Tensor> {
165        self.as_slice().einsum_subscripts(subscripts, backend)
166    }
167}
168
169/// Backend-explicit preallocated-output einsum methods for dtype-erased tensors.
170pub trait TensorEinsumIntoExt {
171    /// Execute an einsum from string notation into caller-provided output.
172    ///
173    /// # Errors
174    ///
175    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
176    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
177    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
178    fn einsum_into<B: TensorBackend>(
179        &self,
180        subscripts: &str,
181        backend: &mut B,
182        out: TensorWrite<'_>,
183    ) -> Result<()>;
184
185    /// Execute an einsum from parsed integer-label subscripts into caller-provided output.
186    ///
187    /// # Errors
188    ///
189    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
190    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
191    /// failure.
192    fn einsum_into_subscripts<B: TensorBackend>(
193        &self,
194        subscripts: &EinsumSubscripts,
195        backend: &mut B,
196        out: TensorWrite<'_>,
197    ) -> Result<()>;
198}
199
200impl TensorEinsumIntoExt for [&Tensor] {
201    fn einsum_into<B: TensorBackend>(
202        &self,
203        subscripts: &str,
204        backend: &mut B,
205        out: TensorWrite<'_>,
206    ) -> Result<()> {
207        let subscripts = parse_subscripts(subscripts, TENSOR_EINSUM_INTO_OP)?;
208        tensor_einsum_into_subscripts(backend, self, &subscripts, out, TENSOR_EINSUM_INTO_OP)
209    }
210
211    fn einsum_into_subscripts<B: TensorBackend>(
212        &self,
213        subscripts: &EinsumSubscripts,
214        backend: &mut B,
215        out: TensorWrite<'_>,
216    ) -> Result<()> {
217        let subscripts = Subscripts::from(subscripts);
218        tensor_einsum_into_subscripts(backend, self, &subscripts, out, TENSOR_EINSUM_INTO_OP)
219    }
220}
221
222impl<const N: usize> TensorEinsumIntoExt for [&Tensor; N] {
223    fn einsum_into<B: TensorBackend>(
224        &self,
225        subscripts: &str,
226        backend: &mut B,
227        out: TensorWrite<'_>,
228    ) -> Result<()> {
229        self.as_slice().einsum_into(subscripts, backend, out)
230    }
231
232    fn einsum_into_subscripts<B: TensorBackend>(
233        &self,
234        subscripts: &EinsumSubscripts,
235        backend: &mut B,
236        out: TensorWrite<'_>,
237    ) -> Result<()> {
238        self.as_slice()
239            .einsum_into_subscripts(subscripts, backend, out)
240    }
241}
242
243/// Backend-explicit einsum methods for typed concrete tensors.
244///
245/// The result keeps the same scalar type as the inputs. Mixed dtypes should use
246/// [`TensorEinsumExt`] on dtype-erased [`Tensor`] values instead.
247///
248/// # Examples
249///
250/// ```
251/// use tenferro_cpu::CpuBackend;
252/// use tenferro_einsum::TypedTensorEinsumExt;
253/// use tenferro_tensor::TypedTensor;
254///
255/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
256/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![3, 4], vec![1.0; 12]).unwrap();
257/// let mut backend = CpuBackend::new();
258///
259/// let out = [&lhs, &rhs].einsum("ij,jk->ik", &mut backend)?;
260/// assert_eq!(out.shape(), &[2, 4]);
261/// # Ok::<(), tenferro_einsum::Error>(())
262/// ```
263pub trait TypedTensorEinsumExt<T: TensorScalar> {
264    /// Execute an einsum from string notation.
265    ///
266    /// # Errors
267    ///
268    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
269    /// [`Error::Validation`] with shape or rank payloads for an invalid
270    /// contraction, or [`Error::Tensor`] for a typed backend failure.
271    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B)
272        -> Result<TypedTensor<T>>;
273
274    /// Execute an einsum from parsed integer-label subscripts.
275    ///
276    /// # Errors
277    ///
278    /// Returns [`Error::Validation`] with shape or rank payloads for an invalid
279    /// contraction, or [`Error::Tensor`] for a typed backend failure.
280    fn einsum_subscripts<B: TensorBackend>(
281        &self,
282        subscripts: &EinsumSubscripts,
283        backend: &mut B,
284    ) -> Result<TypedTensor<T>>;
285}
286
287impl<T: TensorScalar> TypedTensorEinsumExt<T> for [&TypedTensor<T>] {
288    fn einsum<B: TensorBackend>(
289        &self,
290        subscripts: &str,
291        backend: &mut B,
292    ) -> Result<TypedTensor<T>> {
293        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_EINSUM_OP)?;
294        typed_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
295    }
296
297    fn einsum_subscripts<B: TensorBackend>(
298        &self,
299        subscripts: &EinsumSubscripts,
300        backend: &mut B,
301    ) -> Result<TypedTensor<T>> {
302        let subscripts = Subscripts::from(subscripts);
303        typed_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
304    }
305}
306
307impl<T: TensorScalar, const N: usize> TypedTensorEinsumExt<T> for [&TypedTensor<T>; N] {
308    fn einsum<B: TensorBackend>(
309        &self,
310        subscripts: &str,
311        backend: &mut B,
312    ) -> Result<TypedTensor<T>> {
313        self.as_slice().einsum(subscripts, backend)
314    }
315
316    fn einsum_subscripts<B: TensorBackend>(
317        &self,
318        subscripts: &EinsumSubscripts,
319        backend: &mut B,
320    ) -> Result<TypedTensor<T>> {
321        self.as_slice().einsum_subscripts(subscripts, backend)
322    }
323}
324
325/// Backend-explicit einsum methods for typed borrowed views.
326///
327/// The `_read` suffix distinguishes this borrowed-view surface from
328/// [`TypedTensorEinsumExt`], whose unsuffixed methods accept only owned compact
329/// typed tensors.
330///
331/// # Examples
332///
333/// ```
334/// use tenferro_cpu::CpuBackend;
335/// use tenferro_einsum::TypedTensorReadEinsumExt;
336/// use tenferro_tensor::TypedTensor;
337///
338/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 2.0])?;
339/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![3.0, 4.0])?;
340/// let mut backend = CpuBackend::new();
341/// let result = [lhs.as_view(), rhs.as_view()].einsum_read("i,i->", &mut backend)?;
342/// assert_eq!(result.as_slice()?, &[11.0]);
343/// # Ok::<(), tenferro_einsum::Error>(())
344/// ```
345pub trait TypedTensorReadEinsumExt<T: TensorScalar> {
346    /// Execute an einsum from string notation over typed borrowed views.
347    ///
348    /// # Examples
349    ///
350    /// ```
351    /// use tenferro_cpu::CpuBackend;
352    /// use tenferro_einsum::TypedTensorReadEinsumExt;
353    /// use tenferro_tensor::TypedTensor;
354    ///
355    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
356    /// let mut backend = CpuBackend::new();
357    /// let result = [input.as_view()].einsum_read("i->i", &mut backend)?;
358    /// assert_eq!(result.as_slice()?, &[2.0, 3.0]);
359    /// # Ok::<(), tenferro_einsum::Error>(())
360    /// ```
361    ///
362    /// # Errors
363    ///
364    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
365    /// [`Error::Validation`] with shape or rank payloads for incompatible
366    /// views, or [`Error::Tensor`] for a typed backend failure.
367    fn einsum_read<B: TensorBackend>(
368        &self,
369        subscripts: &str,
370        backend: &mut B,
371    ) -> Result<TypedTensor<T>>;
372
373    /// Execute an einsum from parsed integer-label subscripts over typed
374    /// borrowed views.
375    ///
376    /// # Examples
377    ///
378    /// ```
379    /// use tenferro_cpu::CpuBackend;
380    /// use tenferro_einsum::{EinsumSubscripts, TypedTensorReadEinsumExt};
381    /// use tenferro_tensor::TypedTensor;
382    ///
383    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
384    /// let subscripts = EinsumSubscripts::new(&[&[0]], &[0]);
385    /// let mut backend = CpuBackend::new();
386    /// let result = [input.as_view()].einsum_read_subscripts(&subscripts, &mut backend)?;
387    /// assert_eq!(result.as_slice()?, &[2.0, 3.0]);
388    /// # Ok::<(), tenferro_einsum::Error>(())
389    /// ```
390    ///
391    /// # Errors
392    ///
393    /// Returns [`Error::Validation`] with shape or rank payloads for
394    /// incompatible views, or [`Error::Tensor`] for a typed backend failure.
395    fn einsum_read_subscripts<B: TensorBackend>(
396        &self,
397        subscripts: &EinsumSubscripts,
398        backend: &mut B,
399    ) -> Result<TypedTensor<T>>;
400}
401
402impl<'a, T: TensorScalar> TypedTensorReadEinsumExt<T> for [TypedTensorView<'a, T>] {
403    fn einsum_read<B: TensorBackend>(
404        &self,
405        subscripts: &str,
406        backend: &mut B,
407    ) -> Result<TypedTensor<T>> {
408        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_READ_EINSUM_OP)?;
409        typed_view_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_READ_EINSUM_OP)
410    }
411
412    fn einsum_read_subscripts<B: TensorBackend>(
413        &self,
414        subscripts: &EinsumSubscripts,
415        backend: &mut B,
416    ) -> Result<TypedTensor<T>> {
417        let subscripts = Subscripts::from(subscripts);
418        typed_view_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_READ_EINSUM_OP)
419    }
420}
421
422impl<'a, T: TensorScalar, const N: usize> TypedTensorReadEinsumExt<T>
423    for [TypedTensorView<'a, T>; N]
424{
425    fn einsum_read<B: TensorBackend>(
426        &self,
427        subscripts: &str,
428        backend: &mut B,
429    ) -> Result<TypedTensor<T>> {
430        self.as_slice().einsum_read(subscripts, backend)
431    }
432
433    fn einsum_read_subscripts<B: TensorBackend>(
434        &self,
435        subscripts: &EinsumSubscripts,
436        backend: &mut B,
437    ) -> Result<TypedTensor<T>> {
438        self.as_slice().einsum_read_subscripts(subscripts, backend)
439    }
440}
441
442/// Backend-explicit preallocated-output einsum methods for typed concrete tensors.
443pub trait TypedTensorEinsumIntoExt<T: TensorScalar> {
444    /// Execute an einsum from string notation into caller-provided typed output.
445    ///
446    /// # Errors
447    ///
448    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
449    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
450    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
451    fn einsum_into<'out, B, O>(&self, subscripts: &str, backend: &mut B, out: O) -> Result<()>
452    where
453        B: TensorBackend,
454        O: Into<TypedTensorWrite<'out, T>>;
455
456    /// Execute an einsum from parsed integer-label subscripts into caller-provided typed output.
457    ///
458    /// # Errors
459    ///
460    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
461    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
462    /// failure.
463    fn einsum_into_subscripts<'out, B, O>(
464        &self,
465        subscripts: &EinsumSubscripts,
466        backend: &mut B,
467        out: O,
468    ) -> Result<()>
469    where
470        B: TensorBackend,
471        O: Into<TypedTensorWrite<'out, T>>;
472}
473
474impl<T: TensorScalar> TypedTensorEinsumIntoExt<T> for [&TypedTensor<T>] {
475    fn einsum_into<'out, B, O>(&self, subscripts: &str, backend: &mut B, out: O) -> Result<()>
476    where
477        B: TensorBackend,
478        O: Into<TypedTensorWrite<'out, T>>,
479    {
480        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_EINSUM_INTO_OP)?;
481        typed_einsum_into_subscripts(
482            backend,
483            self,
484            &subscripts,
485            out.into(),
486            TYPED_TENSOR_EINSUM_INTO_OP,
487        )
488    }
489
490    fn einsum_into_subscripts<'out, B, O>(
491        &self,
492        subscripts: &EinsumSubscripts,
493        backend: &mut B,
494        out: O,
495    ) -> Result<()>
496    where
497        B: TensorBackend,
498        O: Into<TypedTensorWrite<'out, T>>,
499    {
500        let subscripts = Subscripts::from(subscripts);
501        typed_einsum_into_subscripts(
502            backend,
503            self,
504            &subscripts,
505            out.into(),
506            TYPED_TENSOR_EINSUM_INTO_OP,
507        )
508    }
509}
510
511impl<T: TensorScalar, const N: usize> TypedTensorEinsumIntoExt<T> for [&TypedTensor<T>; N] {
512    fn einsum_into<'out, B, O>(&self, subscripts: &str, backend: &mut B, out: O) -> Result<()>
513    where
514        B: TensorBackend,
515        O: Into<TypedTensorWrite<'out, T>>,
516    {
517        self.as_slice().einsum_into(subscripts, backend, out)
518    }
519
520    fn einsum_into_subscripts<'out, B, O>(
521        &self,
522        subscripts: &EinsumSubscripts,
523        backend: &mut B,
524        out: O,
525    ) -> Result<()>
526    where
527        B: TensorBackend,
528        O: Into<TypedTensorWrite<'out, T>>,
529    {
530        self.as_slice()
531            .einsum_into_subscripts(subscripts, backend, out)
532    }
533}
534
535/// Backend-explicit preallocated-output einsum methods for typed borrowed views.
536///
537/// # Examples
538///
539/// ```
540/// use tenferro_cpu::CpuBackend;
541/// use tenferro_einsum::TypedTensorReadEinsumIntoExt;
542/// use tenferro_tensor::TypedTensor;
543///
544/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 2.0])?;
545/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![3.0, 4.0])?;
546/// let mut output = TypedTensor::<f64>::from_vec_col_major(vec![], vec![0.0])?;
547/// let mut backend = CpuBackend::new();
548/// [lhs.as_view(), rhs.as_view()].einsum_read_into("i,i->", &mut backend, &mut output)?;
549/// assert_eq!(output.as_slice()?, &[11.0]);
550/// # Ok::<(), tenferro_einsum::Error>(())
551/// ```
552pub trait TypedTensorReadEinsumIntoExt<T: TensorScalar> {
553    /// Execute an einsum from string notation over typed borrowed views into a
554    /// caller-provided typed output.
555    ///
556    /// # Examples
557    ///
558    /// ```
559    /// use tenferro_cpu::CpuBackend;
560    /// use tenferro_einsum::TypedTensorReadEinsumIntoExt;
561    /// use tenferro_tensor::TypedTensor;
562    ///
563    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
564    /// let mut output = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![0.0; 2])?;
565    /// let mut backend = CpuBackend::new();
566    /// [input.as_view()].einsum_read_into("i->i", &mut backend, &mut output)?;
567    /// assert_eq!(output.as_slice()?, &[2.0, 3.0]);
568    /// # Ok::<(), tenferro_einsum::Error>(())
569    /// ```
570    ///
571    /// # Errors
572    ///
573    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
574    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
575    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
576    fn einsum_read_into<'out, B, O>(&self, subscripts: &str, backend: &mut B, out: O) -> Result<()>
577    where
578        B: TensorBackend,
579        O: Into<TypedTensorWrite<'out, T>>;
580
581    /// Execute an einsum from parsed integer-label subscripts over typed
582    /// borrowed views into a caller-provided typed output.
583    ///
584    /// # Examples
585    ///
586    /// ```
587    /// use tenferro_cpu::CpuBackend;
588    /// use tenferro_einsum::{EinsumSubscripts, TypedTensorReadEinsumIntoExt};
589    /// use tenferro_tensor::TypedTensor;
590    ///
591    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
592    /// let mut output = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![0.0; 2])?;
593    /// let subscripts = EinsumSubscripts::new(&[&[0]], &[0]);
594    /// let mut backend = CpuBackend::new();
595    /// [input.as_view()].einsum_read_into_subscripts(&subscripts, &mut backend, &mut output)?;
596    /// assert_eq!(output.as_slice()?, &[2.0, 3.0]);
597    /// # Ok::<(), tenferro_einsum::Error>(())
598    /// ```
599    ///
600    /// # Errors
601    ///
602    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
603    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
604    /// failure.
605    fn einsum_read_into_subscripts<'out, B, O>(
606        &self,
607        subscripts: &EinsumSubscripts,
608        backend: &mut B,
609        out: O,
610    ) -> Result<()>
611    where
612        B: TensorBackend,
613        O: Into<TypedTensorWrite<'out, T>>;
614}
615
616impl<'a, T: TensorScalar> TypedTensorReadEinsumIntoExt<T> for [TypedTensorView<'a, T>] {
617    fn einsum_read_into<'out, B, O>(&self, subscripts: &str, backend: &mut B, out: O) -> Result<()>
618    where
619        B: TensorBackend,
620        O: Into<TypedTensorWrite<'out, T>>,
621    {
622        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_READ_EINSUM_INTO_OP)?;
623        typed_view_einsum_into_subscripts(
624            backend,
625            self,
626            &subscripts,
627            out.into(),
628            TYPED_TENSOR_READ_EINSUM_INTO_OP,
629        )
630    }
631
632    fn einsum_read_into_subscripts<'out, B, O>(
633        &self,
634        subscripts: &EinsumSubscripts,
635        backend: &mut B,
636        out: O,
637    ) -> Result<()>
638    where
639        B: TensorBackend,
640        O: Into<TypedTensorWrite<'out, T>>,
641    {
642        let subscripts = Subscripts::from(subscripts);
643        typed_view_einsum_into_subscripts(
644            backend,
645            self,
646            &subscripts,
647            out.into(),
648            TYPED_TENSOR_READ_EINSUM_INTO_OP,
649        )
650    }
651}
652
653impl<'a, T: TensorScalar, const N: usize> TypedTensorReadEinsumIntoExt<T>
654    for [TypedTensorView<'a, T>; N]
655{
656    fn einsum_read_into<'out, B, O>(&self, subscripts: &str, backend: &mut B, out: O) -> Result<()>
657    where
658        B: TensorBackend,
659        O: Into<TypedTensorWrite<'out, T>>,
660    {
661        self.as_slice().einsum_read_into(subscripts, backend, out)
662    }
663
664    fn einsum_read_into_subscripts<'out, B, O>(
665        &self,
666        subscripts: &EinsumSubscripts,
667        backend: &mut B,
668        out: O,
669    ) -> Result<()>
670    where
671        B: TensorBackend,
672        O: Into<TypedTensorWrite<'out, T>>,
673    {
674        self.as_slice()
675            .einsum_read_into_subscripts(subscripts, backend, out)
676    }
677}
678
679/// Backend-explicit einsum methods for [`TensorRead`] inputs.
680///
681/// Use this surface when an input is a borrowed tensor view rather than an
682/// owned compact [`Tensor`]. The `_read` suffix follows the repository-wide
683/// convention for APIs that explicitly accept read-oriented borrowed inputs.
684///
685/// # Examples
686///
687/// ```
688/// use tenferro_cpu::CpuBackend;
689/// use tenferro_einsum::TensorReadEinsumExt;
690/// use tenferro_tensor::{Tensor, TensorRead, TensorView};
691///
692/// let shape = [2, 3];
693/// let data = [1.0_f64; 6];
694/// let rhs = Tensor::from_vec_col_major(vec![3], vec![1.0_f64; 3]).unwrap();
695/// let inputs = [
696///     TensorRead::from_view(TensorView::f64(&shape, &data)?),
697///     TensorRead::from_tensor(&rhs),
698/// ];
699/// let mut backend = CpuBackend::new();
700///
701/// let out = inputs.einsum_read("ij,j->i", &mut backend)?;
702/// assert_eq!(out.shape(), &[2]);
703/// # Ok::<(), tenferro_einsum::Error>(())
704/// ```
705pub trait TensorReadEinsumExt {
706    /// Execute an einsum from string notation over read-only tensor inputs.
707    ///
708    /// # Errors
709    ///
710    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
711    /// [`Error::Validation`] with shape, rank, or dtype payloads for an invalid
712    /// contraction, or [`Error::Tensor`] for a typed backend failure.
713    fn einsum_read<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor>;
714
715    /// Execute an einsum from parsed integer-label subscripts over read-only
716    /// tensor inputs.
717    ///
718    /// # Errors
719    ///
720    /// Returns [`Error::Validation`] with shape, rank, or dtype payloads for an
721    /// invalid contraction, or [`Error::Tensor`] for a typed backend failure.
722    fn einsum_read_subscripts<B: TensorBackend>(
723        &self,
724        subscripts: &EinsumSubscripts,
725        backend: &mut B,
726    ) -> Result<Tensor>;
727}
728
729impl<'a> TensorReadEinsumExt for [TensorRead<'a>] {
730    fn einsum_read<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
731        let subscripts = parse_subscripts(subscripts, TENSOR_READ_EINSUM_OP)?;
732        eager_einsum_read_subscripts(backend, self, &subscripts).map_err(Error::from)
733    }
734
735    fn einsum_read_subscripts<B: TensorBackend>(
736        &self,
737        subscripts: &EinsumSubscripts,
738        backend: &mut B,
739    ) -> Result<Tensor> {
740        let subscripts = Subscripts::from(subscripts);
741        eager_einsum_read_subscripts(backend, self, &subscripts).map_err(Error::from)
742    }
743}
744
745impl<'a, const N: usize> TensorReadEinsumExt for [TensorRead<'a>; N] {
746    fn einsum_read<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
747        self.as_slice().einsum_read(subscripts, backend)
748    }
749
750    fn einsum_read_subscripts<B: TensorBackend>(
751        &self,
752        subscripts: &EinsumSubscripts,
753        backend: &mut B,
754    ) -> Result<Tensor> {
755        self.as_slice().einsum_read_subscripts(subscripts, backend)
756    }
757}
758
759/// Backend-explicit preallocated-output einsum methods for [`TensorRead`] inputs.
760pub trait TensorReadEinsumIntoExt {
761    /// Execute an einsum from string notation over read-only inputs into caller-provided output.
762    ///
763    /// # Errors
764    ///
765    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
766    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
767    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
768    fn einsum_read_into<B: TensorBackend>(
769        &self,
770        subscripts: &str,
771        backend: &mut B,
772        out: TensorWrite<'_>,
773    ) -> Result<()>;
774
775    /// Execute an einsum from parsed integer-label subscripts over read-only inputs into output.
776    ///
777    /// # Errors
778    ///
779    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
780    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
781    /// failure.
782    fn einsum_read_into_subscripts<B: TensorBackend>(
783        &self,
784        subscripts: &EinsumSubscripts,
785        backend: &mut B,
786        out: TensorWrite<'_>,
787    ) -> Result<()>;
788}
789
790impl<'a> TensorReadEinsumIntoExt for [TensorRead<'a>] {
791    fn einsum_read_into<B: TensorBackend>(
792        &self,
793        subscripts: &str,
794        backend: &mut B,
795        out: TensorWrite<'_>,
796    ) -> Result<()> {
797        let subscripts = parse_subscripts(subscripts, TENSOR_READ_EINSUM_INTO_OP)?;
798        tensor_read_einsum_into_subscripts(
799            backend,
800            self,
801            &subscripts,
802            out,
803            TENSOR_READ_EINSUM_INTO_OP,
804        )
805    }
806
807    fn einsum_read_into_subscripts<B: TensorBackend>(
808        &self,
809        subscripts: &EinsumSubscripts,
810        backend: &mut B,
811        out: TensorWrite<'_>,
812    ) -> Result<()> {
813        let subscripts = Subscripts::from(subscripts);
814        tensor_read_einsum_into_subscripts(
815            backend,
816            self,
817            &subscripts,
818            out,
819            TENSOR_READ_EINSUM_INTO_OP,
820        )
821    }
822}
823
824impl<'a, const N: usize> TensorReadEinsumIntoExt for [TensorRead<'a>; N] {
825    fn einsum_read_into<B: TensorBackend>(
826        &self,
827        subscripts: &str,
828        backend: &mut B,
829        out: TensorWrite<'_>,
830    ) -> Result<()> {
831        self.as_slice().einsum_read_into(subscripts, backend, out)
832    }
833
834    fn einsum_read_into_subscripts<B: TensorBackend>(
835        &self,
836        subscripts: &EinsumSubscripts,
837        backend: &mut B,
838        out: TensorWrite<'_>,
839    ) -> Result<()> {
840        self.as_slice()
841            .einsum_read_into_subscripts(subscripts, backend, out)
842    }
843}
844
845/// Prepared concrete einsum plan for repeated executions with fixed input
846/// dtype and shape metadata.
847///
848/// Preparing a plan parses and optimizes the contraction tree once. Execution
849/// validates the later inputs against the prepared dtype and shape contract,
850/// then runs the stored tree without re-planning.
851///
852/// # Examples
853///
854/// ```
855/// use tenferro_cpu::CpuBackend;
856/// use tenferro_einsum::ConcreteEinsumPlan;
857/// use tenferro_tensor::Tensor;
858///
859/// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
860/// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
861/// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik")?;
862///
863/// let mut backend = CpuBackend::new();
864/// let out = plan.execute([&lhs, &rhs], &mut backend)?;
865/// assert_eq!(out.shape(), &[2, 4]);
866/// # Ok::<(), tenferro_einsum::Error>(())
867/// ```
868#[derive(Debug)]
869pub struct ConcreteEinsumPlan {
870    tree: ContractionTree,
871    inputs: Vec<ConcreteEinsumInputSpec>,
872}
873
874impl ConcreteEinsumPlan {
875    /// Prepare a plan from dtype-erased concrete tensor inputs and string
876    /// notation.
877    ///
878    /// # Errors
879    ///
880    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
881    /// [`Error::Validation`] for rank, shape, or dtype contract violations, or
882    /// [`Error::Planning`] when no valid contraction tree can be built.
883    pub fn prepare<'a, I>(inputs: I, subscripts: &str) -> Result<Self>
884    where
885        I: AsRef<[&'a Tensor]>,
886    {
887        let subscripts = parse_subscripts(subscripts, PLAN_PREPARE_OP)?;
888        Self::prepare_subscripts_internal(input_specs(inputs.as_ref()), &subscripts)
889    }
890
891    /// Prepare a plan from dtype-erased concrete tensor inputs and parsed
892    /// integer-label subscripts.
893    ///
894    /// # Errors
895    ///
896    /// Returns [`Error::Validation`] for rank, shape, or dtype contract
897    /// violations, or [`Error::Planning`] when no valid contraction tree can be
898    /// built.
899    pub fn prepare_subscripts<'a, I>(inputs: I, subscripts: &EinsumSubscripts) -> Result<Self>
900    where
901        I: AsRef<[&'a Tensor]>,
902    {
903        let subscripts = Subscripts::from(subscripts);
904        Self::prepare_subscripts_internal(input_specs(inputs.as_ref()), &subscripts)
905    }
906
907    /// Prepare a plan from typed concrete tensor inputs and string notation.
908    ///
909    /// # Errors
910    ///
911    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
912    /// [`Error::Validation`] for rank or shape contract violations, or
913    /// [`Error::Planning`] when no valid contraction tree can be built.
914    pub fn prepare_typed<'a, T, I>(inputs: I, subscripts: &str) -> Result<Self>
915    where
916        T: TensorScalar,
917        I: AsRef<[&'a TypedTensor<T>]>,
918    {
919        let subscripts = parse_subscripts(subscripts, PLAN_PREPARE_OP)?;
920        Self::prepare_subscripts_internal(typed_input_specs(inputs.as_ref()), &subscripts)
921    }
922
923    /// Prepare a plan from typed concrete tensor inputs and parsed integer-label
924    /// subscripts.
925    ///
926    /// # Errors
927    ///
928    /// Returns [`Error::Validation`] for rank or shape contract violations, or
929    /// [`Error::Planning`] when no valid contraction tree can be built.
930    pub fn prepare_typed_subscripts<'a, T, I>(
931        inputs: I,
932        subscripts: &EinsumSubscripts,
933    ) -> Result<Self>
934    where
935        T: TensorScalar,
936        I: AsRef<[&'a TypedTensor<T>]>,
937    {
938        let subscripts = Subscripts::from(subscripts);
939        Self::prepare_subscripts_internal(typed_input_specs(inputs.as_ref()), &subscripts)
940    }
941
942    /// Prepare a plan from read-only tensor inputs and string notation.
943    ///
944    /// # Errors
945    ///
946    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
947    /// [`Error::Validation`] for rank, shape, or dtype contract violations, or
948    /// [`Error::Planning`] when no valid contraction tree can be built.
949    pub fn prepare_read<'a, I>(inputs: I, subscripts: &str) -> Result<Self>
950    where
951        I: AsRef<[TensorRead<'a>]>,
952    {
953        let subscripts = parse_subscripts(subscripts, PLAN_PREPARE_OP)?;
954        Self::prepare_subscripts_internal(read_input_specs(inputs.as_ref()), &subscripts)
955    }
956
957    /// Prepare a plan from read-only tensor inputs and parsed integer-label
958    /// subscripts.
959    ///
960    /// # Errors
961    ///
962    /// Returns [`Error::Validation`] for rank, shape, or dtype contract
963    /// violations, or [`Error::Planning`] when no valid contraction tree can be
964    /// built.
965    pub fn prepare_read_subscripts<'a, I>(inputs: I, subscripts: &EinsumSubscripts) -> Result<Self>
966    where
967        I: AsRef<[TensorRead<'a>]>,
968    {
969        let subscripts = Subscripts::from(subscripts);
970        Self::prepare_subscripts_internal(read_input_specs(inputs.as_ref()), &subscripts)
971    }
972
973    /// Execute this plan on dtype-erased concrete tensor inputs.
974    ///
975    /// # Errors
976    ///
977    /// Returns [`Error::Validation`] when inputs differ from the prepared
978    /// rank, shape, or dtype contract, or [`Error::Tensor`] for a typed backend
979    /// failure.
980    pub fn execute<'a, I, B>(&self, inputs: I, backend: &mut B) -> Result<Tensor>
981    where
982        I: AsRef<[&'a Tensor]>,
983        B: TensorBackend,
984    {
985        let inputs = inputs.as_ref();
986        self.validate_inputs(&input_specs(inputs), PLAN_EXECUTE_OP)?;
987        backend
988            .with_backend_session(|exec| eager_einsum_exec(exec, inputs, &self.tree))
989            .map_err(Error::from)
990    }
991
992    /// Execute this plan on typed concrete tensor inputs.
993    ///
994    /// # Errors
995    ///
996    /// Returns [`Error::Validation`] when inputs differ from the prepared rank
997    /// or shape contract, or [`Error::Tensor`] for a typed backend failure.
998    pub fn execute_typed<'a, T, I, B>(&self, inputs: I, backend: &mut B) -> Result<TypedTensor<T>>
999    where
1000        T: TensorScalar,
1001        I: AsRef<[&'a TypedTensor<T>]>,
1002        B: TensorBackend,
1003    {
1004        let inputs = inputs.as_ref();
1005        self.validate_inputs(&typed_input_specs(inputs), PLAN_EXECUTE_OP)?;
1006        let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1007        let result = backend
1008            .with_backend_session(|exec| eager_einsum_exec_read(exec, &reads, &self.tree))?;
1009        into_typed_result(result, PLAN_EXECUTE_OP)
1010    }
1011
1012    /// Execute this plan on read-only tensor inputs.
1013    ///
1014    /// # Errors
1015    ///
1016    /// Returns [`Error::Validation`] when inputs differ from the prepared
1017    /// rank, shape, or dtype contract, or [`Error::Tensor`] for a typed backend
1018    /// failure.
1019    pub fn execute_read<'a, I, B>(&self, inputs: I, backend: &mut B) -> Result<Tensor>
1020    where
1021        I: AsRef<[TensorRead<'a>]>,
1022        B: TensorBackend,
1023    {
1024        let inputs = inputs.as_ref();
1025        self.validate_inputs(&read_input_specs(inputs), PLAN_EXECUTE_OP)?;
1026        backend
1027            .with_backend_session(|exec| eager_einsum_exec_read(exec, inputs, &self.tree))
1028            .map_err(Error::from)
1029    }
1030
1031    /// Execute this plan on dtype-erased concrete tensor inputs into caller-provided output.
1032    ///
1033    /// # Errors
1034    ///
1035    /// Returns [`Error::Validation`] for input or output rank, shape, or dtype
1036    /// mismatches, or [`Error::Tensor`] for a typed backend failure.
1037    pub fn execute_into<'a, I, B>(
1038        &self,
1039        inputs: I,
1040        backend: &mut B,
1041        out: TensorWrite<'_>,
1042    ) -> Result<()>
1043    where
1044        I: AsRef<[&'a Tensor]>,
1045        B: TensorBackend,
1046    {
1047        let inputs = inputs.as_ref();
1048        let specs = input_specs(inputs);
1049        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1050        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1051        let reads: Vec<_> = inputs
1052            .iter()
1053            .map(|tensor| TensorRead::from_tensor(tensor))
1054            .collect();
1055        backend
1056            .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &self.tree, out))
1057            .map_err(Error::from)
1058    }
1059
1060    /// Execute this plan on typed concrete tensor inputs into caller-provided output.
1061    ///
1062    /// # Errors
1063    ///
1064    /// Returns [`Error::Validation`] for input or output rank or shape
1065    /// mismatches, or [`Error::Tensor`] for a typed backend failure.
1066    pub fn execute_typed_into<'a, 'out, T, I, B, O>(
1067        &self,
1068        inputs: I,
1069        backend: &mut B,
1070        out: O,
1071    ) -> Result<()>
1072    where
1073        T: TensorScalar,
1074        I: AsRef<[&'a TypedTensor<T>]>,
1075        B: TensorBackend,
1076        O: Into<TypedTensorWrite<'out, T>>,
1077    {
1078        let inputs = inputs.as_ref();
1079        let specs = typed_input_specs(inputs);
1080        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1081        let out = out.into().into_tensor_write();
1082        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1083        let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1084        backend
1085            .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &self.tree, out))
1086            .map_err(Error::from)
1087    }
1088
1089    /// Execute this plan on read-only tensor inputs into caller-provided output.
1090    ///
1091    /// # Errors
1092    ///
1093    /// Returns [`Error::Validation`] for input or output rank, shape, or dtype
1094    /// mismatches, or [`Error::Tensor`] for a typed backend failure.
1095    pub fn execute_read_into<'a, I, B>(
1096        &self,
1097        inputs: I,
1098        backend: &mut B,
1099        out: TensorWrite<'_>,
1100    ) -> Result<()>
1101    where
1102        I: AsRef<[TensorRead<'a>]>,
1103        B: TensorBackend,
1104    {
1105        let inputs = inputs.as_ref();
1106        let specs = read_input_specs(inputs);
1107        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1108        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1109        backend
1110            .with_backend_session(|exec| eager_einsum_exec_read_into(exec, inputs, &self.tree, out))
1111            .map_err(Error::from)
1112    }
1113
1114    /// Execute this plan on read-only inputs with scaled output accumulation.
1115    ///
1116    /// `accumulation` follows the dot-general contract:
1117    /// `out = alpha * einsum(inputs) + beta * out`.
1118    ///
1119    /// # Examples
1120    ///
1121    /// ```rust
1122    /// use tenferro_cpu::CpuBackend;
1123    /// use tenferro_einsum::ConcreteEinsumPlan;
1124    /// use tenferro_tensor::{
1125    ///     DotGeneralAccumulation, DType, Tensor, TensorRead, TensorWrite,
1126    /// };
1127    ///
1128    /// let lhs = Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?;
1129    /// let rhs = Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?;
1130    /// let mut out = Tensor::from_vec_col_major(vec![], vec![1.0_f64])?;
1131    /// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "i,i->")?;
1132    /// let mut backend = CpuBackend::new();
1133    /// plan.execute_read_into_accum(
1134    ///     [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
1135    ///     &mut backend,
1136    ///     DotGeneralAccumulation::add_to(DType::F64)?,
1137    ///     TensorWrite::from_tensor(&mut out),
1138    /// )?;
1139    /// assert_eq!(out.as_slice::<f64>()?, &[7.0]);
1140    /// # Ok::<(), tenferro_einsum::Error>(())
1141    /// ```
1142    ///
1143    /// # Errors
1144    ///
1145    /// Returns [`Error::Validation`] for input or output rank, shape, or dtype
1146    /// mismatches, [`Error::Numerical`] for an invalid accumulation, or
1147    /// [`Error::Tensor`] for a typed backend failure.
1148    pub fn execute_read_into_accum<'a, I, B>(
1149        &self,
1150        inputs: I,
1151        backend: &mut B,
1152        accumulation: DotGeneralAccumulation,
1153        out: TensorWrite<'_>,
1154    ) -> Result<()>
1155    where
1156        I: AsRef<[TensorRead<'a>]>,
1157        B: TensorBackend,
1158    {
1159        let inputs = inputs.as_ref();
1160        let specs = read_input_specs(inputs);
1161        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1162        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1163        backend
1164            .with_backend_session(|exec| {
1165                eager_einsum_exec_read_into_accum(exec, inputs, &self.tree, accumulation, out)
1166            })
1167            .map_err(Error::from)
1168    }
1169
1170    fn prepare_subscripts_internal(
1171        inputs: Vec<ConcreteEinsumInputSpec>,
1172        subscripts: &Subscripts,
1173    ) -> Result<Self> {
1174        let shapes: Vec<&[usize]> = inputs.iter().map(|input| input.shape.as_slice()).collect();
1175        let tree = plan_subscripts(subscripts, &shapes)?;
1176        Ok(Self { tree, inputs })
1177    }
1178
1179    fn validate_inputs(&self, actual: &[ConcreteEinsumInputSpec], op: &'static str) -> Result<()> {
1180        if actual.len() != self.inputs.len() {
1181            return Err(Error::invalid_argument(
1182                op,
1183                "inputs",
1184                format!(
1185                    "prepared einsum expects {} inputs, got {}",
1186                    self.inputs.len(),
1187                    actual.len()
1188                ),
1189            ));
1190        }
1191
1192        for (expected, actual) in self.inputs.iter().zip(actual.iter()) {
1193            if expected.dtype != actual.dtype {
1194                return Err(Error::dtype_mismatch(op, expected.dtype, actual.dtype));
1195            }
1196            if expected.shape != actual.shape {
1197                return Err(Error::shape_mismatch(
1198                    op,
1199                    expected.shape.clone(),
1200                    actual.shape.clone(),
1201                ));
1202            }
1203        }
1204
1205        Ok(())
1206    }
1207}
1208
1209#[derive(Clone, Debug)]
1210struct ConcreteEinsumInputSpec {
1211    dtype: DType,
1212    shape: Vec<usize>,
1213}
1214
1215fn parse_subscripts(subscripts: &str, _op: &'static str) -> Result<Subscripts> {
1216    Subscripts::parse(subscripts)
1217}
1218
1219fn input_specs(inputs: &[&Tensor]) -> Vec<ConcreteEinsumInputSpec> {
1220    inputs
1221        .iter()
1222        .map(|tensor| ConcreteEinsumInputSpec {
1223            dtype: tensor.dtype(),
1224            shape: tensor.shape().to_vec(),
1225        })
1226        .collect()
1227}
1228
1229fn typed_input_specs<T: TensorScalar>(inputs: &[&TypedTensor<T>]) -> Vec<ConcreteEinsumInputSpec> {
1230    inputs
1231        .iter()
1232        .map(|tensor| ConcreteEinsumInputSpec {
1233            dtype: T::dtype(),
1234            shape: tensor.shape().to_vec(),
1235        })
1236        .collect()
1237}
1238
1239fn typed_view_input_specs<T: TensorScalar>(
1240    inputs: &[TypedTensorView<'_, T>],
1241) -> Vec<ConcreteEinsumInputSpec> {
1242    inputs
1243        .iter()
1244        .map(|tensor| ConcreteEinsumInputSpec {
1245            dtype: T::dtype(),
1246            shape: tensor.shape().to_vec(),
1247        })
1248        .collect()
1249}
1250
1251fn read_input_specs(inputs: &[TensorRead<'_>]) -> Vec<ConcreteEinsumInputSpec> {
1252    inputs
1253        .iter()
1254        .map(|tensor| ConcreteEinsumInputSpec {
1255            dtype: tensor.dtype(),
1256            shape: tensor.shape().to_vec(),
1257        })
1258        .collect()
1259}
1260
1261fn typed_view_einsum_subscripts<T: TensorScalar>(
1262    backend: &mut impl TensorBackend,
1263    inputs: &[TypedTensorView<'_, T>],
1264    subscripts: &Subscripts,
1265    op: &'static str,
1266) -> Result<TypedTensor<T>> {
1267    let reads: Vec<_> = inputs
1268        .iter()
1269        .cloned()
1270        .map(|view| TensorRead::from_view(T::tensor_view(view)))
1271        .collect();
1272    let result = eager_einsum_read_subscripts(backend, &reads, subscripts)?;
1273    into_typed_result(result, op)
1274}
1275
1276fn tensor_einsum_into_subscripts(
1277    backend: &mut impl TensorBackend,
1278    inputs: &[&Tensor],
1279    subscripts: &Subscripts,
1280    out: TensorWrite<'_>,
1281    op: &'static str,
1282) -> Result<()> {
1283    let specs = input_specs(inputs);
1284    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
1285    validate_output(&specs, &plan.tree, &out, op)?;
1286    let reads: Vec<_> = inputs
1287        .iter()
1288        .map(|tensor| TensorRead::from_tensor(tensor))
1289        .collect();
1290    backend
1291        .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &plan.tree, out))
1292        .map_err(Error::from)
1293}
1294
1295fn typed_view_einsum_into_subscripts<T: TensorScalar>(
1296    backend: &mut impl TensorBackend,
1297    inputs: &[TypedTensorView<'_, T>],
1298    subscripts: &Subscripts,
1299    out: TypedTensorWrite<'_, T>,
1300    op: &'static str,
1301) -> Result<()> {
1302    let specs = typed_view_input_specs(inputs);
1303    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
1304    let out = out.into_tensor_write();
1305    validate_output(&specs, &plan.tree, &out, op)?;
1306    let reads: Vec<_> = inputs
1307        .iter()
1308        .cloned()
1309        .map(|view| TensorRead::from_view(T::tensor_view(view)))
1310        .collect();
1311    backend
1312        .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &plan.tree, out))
1313        .map_err(Error::from)
1314}
1315
1316fn typed_einsum_into_subscripts<T: TensorScalar>(
1317    backend: &mut impl TensorBackend,
1318    inputs: &[&TypedTensor<T>],
1319    subscripts: &Subscripts,
1320    out: TypedTensorWrite<'_, T>,
1321    op: &'static str,
1322) -> Result<()> {
1323    let specs = typed_input_specs(inputs);
1324    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
1325    let out = out.into_tensor_write();
1326    validate_output(&specs, &plan.tree, &out, op)?;
1327    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1328    backend
1329        .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &plan.tree, out))
1330        .map_err(Error::from)
1331}
1332
1333fn tensor_read_einsum_into_subscripts(
1334    backend: &mut impl TensorBackend,
1335    inputs: &[TensorRead<'_>],
1336    subscripts: &Subscripts,
1337    out: TensorWrite<'_>,
1338    op: &'static str,
1339) -> Result<()> {
1340    let specs = read_input_specs(inputs);
1341    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
1342    validate_output(&specs, &plan.tree, &out, op)?;
1343    backend
1344        .with_backend_session(|exec| eager_einsum_exec_read_into(exec, inputs, &plan.tree, out))
1345        .map_err(Error::from)
1346}
1347
1348fn validate_output(
1349    inputs: &[ConcreteEinsumInputSpec],
1350    tree: &ContractionTree,
1351    out: &TensorWrite<'_>,
1352    op: &'static str,
1353) -> Result<()> {
1354    let expected = output_spec(inputs, tree, op)?;
1355    if out.dtype() != expected.dtype {
1356        return Err(Error::dtype_mismatch(op, expected.dtype, out.dtype()));
1357    }
1358    if out.shape() != expected.shape.as_slice() {
1359        return Err(Error::shape_mismatch(
1360            op,
1361            out.shape().to_vec(),
1362            expected.shape.clone(),
1363        ));
1364    }
1365    Ok(())
1366}
1367
1368fn output_spec(
1369    inputs: &[ConcreteEinsumInputSpec],
1370    tree: &ContractionTree,
1371    op: &'static str,
1372) -> Result<ConcreteEinsumInputSpec> {
1373    let dtype = inputs
1374        .first()
1375        .ok_or_else(|| {
1376            Error::invalid_argument(op, "inputs", "einsum requires at least one input tensor")
1377        })?
1378        .dtype;
1379    for input in inputs {
1380        if input.dtype != dtype {
1381            return Err(Error::dtype_mismatch(op, dtype, input.dtype));
1382        }
1383    }
1384
1385    let mut output_shape = Vec::with_capacity(tree.subscripts.output.len());
1386    for &label in &tree.subscripts.output {
1387        let mut found = None;
1388        for (input, labels) in inputs.iter().zip(tree.subscripts.inputs.iter()) {
1389            if labels.len() != input.shape.len() {
1390                return Err(Error::rank_mismatch(op, labels.len(), input.shape.len()));
1391            }
1392            if let Some(axis) = labels.iter().position(|candidate| *candidate == label) {
1393                found = Some(input.shape[axis]);
1394                break;
1395            }
1396        }
1397        let Some(extent) = found else {
1398            return Err(Error::invalid_argument(
1399                op,
1400                "output labels",
1401                format!("output label {label} is missing from inputs"),
1402            ));
1403        };
1404        output_shape.push(extent);
1405    }
1406    Ok(ConcreteEinsumInputSpec {
1407        dtype,
1408        shape: output_shape,
1409    })
1410}
1411
1412fn typed_einsum_subscripts<T: TensorScalar>(
1413    backend: &mut impl TensorBackend,
1414    inputs: &[&TypedTensor<T>],
1415    subscripts: &Subscripts,
1416    op: &'static str,
1417) -> Result<TypedTensor<T>> {
1418    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1419    let result = eager_einsum_read_subscripts(backend, &reads, subscripts)?;
1420    into_typed_result(result, op)
1421}
1422
1423pub(crate) fn into_typed_result<T: TensorScalar>(
1424    result: Tensor,
1425    op: &'static str,
1426) -> Result<TypedTensor<T>> {
1427    let actual = result.dtype();
1428    T::into_typed(result).map_err(|_| Error::dtype_mismatch(op, T::dtype(), actual))
1429}
tenferro_einsum/concrete.rs

tenferro_einsum/
concrete.rs
