Skip to main content

tenferro_einsum/
concrete.rs

1//! Public concrete tensor einsum extension API.
2
3use tenferro_tensor::{
4    BackendSession, DType, DotGeneralAccumulation, Tensor, 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_on_session,
11    eager_einsum_subscripts_on_session, plan_subscripts,
12};
13use crate::ellipsis::resolve_einsum_notation;
14use crate::TensorDotAxes;
15use crate::{
16    parse_einsum_notation, ContractionTree, EinsumNotation, EinsumSubscripts, Error, Result,
17    Subscripts,
18};
19
20const TENSOR_EINSUM_INTO_OP: &str = "TensorEinsumIntoExt::einsum_into";
21const TENSOR_READ_EINSUM_INTO_OP: &str = "TensorReadEinsumIntoExt::einsum_read_into";
22const TYPED_TENSOR_EINSUM_OP: &str = "TypedTensorEinsumExt::einsum";
23const TYPED_TENSOR_EINSUM_INTO_OP: &str = "TypedTensorEinsumIntoExt::einsum_into";
24const TYPED_TENSOR_READ_EINSUM_OP: &str = "TypedTensorReadEinsumExt::einsum_read";
25const TYPED_TENSOR_READ_EINSUM_INTO_OP: &str = "TypedTensorReadEinsumIntoExt::einsum_read_into";
26const PLAN_EXECUTE_OP: &str = "ConcreteEinsumPlan::execute";
27const TYPED_TENSOR_TENSORDOT_OP: &str = "TypedTensorTensordotExt::tensordot";
28
29/// Backend-explicit tensordot sugar for dtype-erased concrete tensors.
30pub trait TensorTensordotExt {
31    /// Contract this tensor with `rhs` over explicit axes or an axis count.
32    ///
33    /// # Errors
34    ///
35    /// Returns [`Error::Validation`] with `InvalidArgument`, `RankMismatch`,
36    /// `AxisOutOfBounds`, or `ShapeMismatch` for invalid axes or incompatible
37    /// contracting dimensions, or [`Error::Tensor`] when backend execution
38    /// fails.
39    fn tensordot(
40        &self,
41        rhs: &Tensor,
42        axes: TensorDotAxes<'_>,
43        session: &mut dyn BackendSession,
44    ) -> Result<Tensor>;
45}
46
47impl TensorTensordotExt for Tensor {
48    fn tensordot(
49        &self,
50        rhs: &Tensor,
51        axes: TensorDotAxes<'_>,
52        session: &mut dyn BackendSession,
53    ) -> Result<Tensor> {
54        let config =
55            crate::tensordot::dot_general_config(axes, self.shape().len(), rhs.shape().len())?;
56        crate::tensordot::validate_concrete_contract_dims(self.shape(), rhs.shape(), &config)?;
57        session.dot_general(self, rhs, &config).map_err(Error::from)
58    }
59}
60
61/// Backend-explicit tensordot sugar for typed concrete tensors.
62pub trait TypedTensorTensordotExt<T: TensorScalar> {
63    /// Contract this tensor with `rhs` while preserving its scalar type.
64    ///
65    /// # Errors
66    ///
67    /// Returns [`Error::Validation`] with `InvalidArgument`, `RankMismatch`,
68    /// `AxisOutOfBounds`, or `ShapeMismatch` for invalid axes or incompatible
69    /// contracting dimensions, or [`Error::Tensor`] when backend execution
70    /// fails.
71    fn tensordot(
72        &self,
73        rhs: &TypedTensor<T>,
74        axes: TensorDotAxes<'_>,
75        session: &mut dyn BackendSession,
76    ) -> Result<TypedTensor<T>>;
77}
78
79impl<T: TensorScalar> TypedTensorTensordotExt<T> for TypedTensor<T> {
80    fn tensordot(
81        &self,
82        rhs: &TypedTensor<T>,
83        axes: TensorDotAxes<'_>,
84        session: &mut dyn BackendSession,
85    ) -> Result<TypedTensor<T>> {
86        let config =
87            crate::tensordot::dot_general_config(axes, self.shape().len(), rhs.shape().len())?;
88        crate::tensordot::validate_concrete_contract_dims(self.shape(), rhs.shape(), &config)?;
89        let result = session
90            .dot_general_read(T::tensor_read(self), T::tensor_read(rhs), &config)
91            .map_err(Error::from)?;
92        into_typed_result(result, TYPED_TENSOR_TENSORDOT_OP)
93    }
94}
95
96/// Backend-explicit einsum methods for dtype-erased concrete tensors.
97///
98/// Implementations are provided for slices and fixed-size arrays of
99/// [`Tensor`] references, so both `inputs.as_slice().einsum(...)` and
100/// `[&lhs, &rhs].einsum(...)` work.
101///
102/// # Examples
103///
104/// ```
105/// use tenferro_cpu::CpuBackend;
106/// use tenferro_einsum::TensorEinsumExt;
107/// use tenferro_tensor::{BackendSessionHost, Tensor};
108///
109/// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
110/// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
111/// let mut backend = CpuBackend::new();
112///
113/// let out = backend.with_backend_session(|session| {
114///     [&lhs, &rhs].einsum("ij,jk->ik", session)
115/// })?;
116/// assert_eq!(out.shape(), &[2, 4]);
117/// # Ok::<(), tenferro_einsum::Error>(())
118/// ```
119pub trait TensorEinsumExt {
120    /// Execute an einsum from string notation.
121    ///
122    /// # Errors
123    ///
124    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
125    /// [`Error::Validation`] with shape, rank, or dtype payloads for an invalid
126    /// contraction, or [`Error::Tensor`] for a typed backend failure.
127    fn einsum(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<Tensor>;
128
129    /// Execute an einsum from rank-unresolved string/programmatic notation.
130    ///
131    /// # Examples
132    ///
133    /// ```
134    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
135    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
136    /// assert_eq!(notation.input_count(), 1);
137    /// ```
138    ///
139    /// # Errors
140    ///
141    /// Returns a typed validation or backend error when notation, shapes, or execution are invalid.
142    fn einsum_notation(
143        &self,
144        notation: &EinsumNotation,
145        session: &mut dyn BackendSession,
146    ) -> Result<Tensor>;
147
148    /// Execute an einsum from parsed integer-label subscripts.
149    ///
150    /// # Errors
151    ///
152    /// Returns [`Error::Validation`] with shape, rank, or dtype payloads for an
153    /// invalid contraction, or [`Error::Tensor`] for a typed backend failure.
154    fn einsum_subscripts(
155        &self,
156        subscripts: &EinsumSubscripts,
157        session: &mut dyn BackendSession,
158    ) -> Result<Tensor>;
159}
160
161impl TensorEinsumExt for [&Tensor] {
162    fn einsum(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<Tensor> {
163        let notation = parse_einsum_notation(subscripts)?;
164        self.einsum_notation(&notation, session)
165    }
166
167    fn einsum_notation(
168        &self,
169        notation: &EinsumNotation,
170        session: &mut dyn BackendSession,
171    ) -> Result<Tensor> {
172        let subscripts = resolve_tensor_notation(self, notation)?;
173        eager_einsum_subscripts_on_session(session, self, &subscripts).map_err(Error::from)
174    }
175
176    fn einsum_subscripts(
177        &self,
178        subscripts: &EinsumSubscripts,
179        session: &mut dyn BackendSession,
180    ) -> Result<Tensor> {
181        let subscripts = Subscripts::from(subscripts);
182        eager_einsum_subscripts_on_session(session, self, &subscripts).map_err(Error::from)
183    }
184}
185
186impl<const N: usize> TensorEinsumExt for [&Tensor; N] {
187    fn einsum(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<Tensor> {
188        self.as_slice().einsum(subscripts, session)
189    }
190
191    fn einsum_notation(
192        &self,
193        notation: &EinsumNotation,
194        session: &mut dyn BackendSession,
195    ) -> Result<Tensor> {
196        self.as_slice().einsum_notation(notation, session)
197    }
198
199    fn einsum_subscripts(
200        &self,
201        subscripts: &EinsumSubscripts,
202        session: &mut dyn BackendSession,
203    ) -> Result<Tensor> {
204        self.as_slice().einsum_subscripts(subscripts, session)
205    }
206}
207
208/// Backend-explicit preallocated-output einsum methods for dtype-erased tensors.
209pub trait TensorEinsumIntoExt {
210    /// Execute an einsum from string notation into caller-provided output.
211    ///
212    /// # Errors
213    ///
214    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
215    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
216    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
217    fn einsum_into(
218        &self,
219        subscripts: &str,
220        session: &mut dyn BackendSession,
221        out: TensorWrite<'_>,
222    ) -> Result<()>;
223
224    /// Execute an einsum from rank-unresolved notation into caller-provided output.
225    ///
226    /// # Examples
227    ///
228    /// ```
229    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
230    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
231    /// assert_eq!(notation.input_count(), 1);
232    /// ```
233    ///
234    /// # Errors
235    ///
236    /// Returns a typed validation or backend error when notation, shapes, or output are invalid.
237    fn einsum_into_notation(
238        &self,
239        notation: &EinsumNotation,
240        session: &mut dyn BackendSession,
241        out: TensorWrite<'_>,
242    ) -> Result<()>;
243
244    /// Execute an einsum from parsed integer-label subscripts into caller-provided output.
245    ///
246    /// # Errors
247    ///
248    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
249    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
250    /// failure.
251    fn einsum_into_subscripts(
252        &self,
253        subscripts: &EinsumSubscripts,
254        session: &mut dyn BackendSession,
255        out: TensorWrite<'_>,
256    ) -> Result<()>;
257}
258
259impl TensorEinsumIntoExt for [&Tensor] {
260    fn einsum_into(
261        &self,
262        subscripts: &str,
263        session: &mut dyn BackendSession,
264        out: TensorWrite<'_>,
265    ) -> Result<()> {
266        let notation = parse_einsum_notation(subscripts)?;
267        self.einsum_into_notation(&notation, session, out)
268    }
269
270    fn einsum_into_notation(
271        &self,
272        notation: &EinsumNotation,
273        session: &mut dyn BackendSession,
274        out: TensorWrite<'_>,
275    ) -> Result<()> {
276        let subscripts = resolve_tensor_notation(self, notation)?;
277        tensor_einsum_into_subscripts(session, self, &subscripts, out, TENSOR_EINSUM_INTO_OP)
278    }
279
280    fn einsum_into_subscripts(
281        &self,
282        subscripts: &EinsumSubscripts,
283        session: &mut dyn BackendSession,
284        out: TensorWrite<'_>,
285    ) -> Result<()> {
286        let subscripts = Subscripts::from(subscripts);
287        tensor_einsum_into_subscripts(session, self, &subscripts, out, TENSOR_EINSUM_INTO_OP)
288    }
289}
290
291impl<const N: usize> TensorEinsumIntoExt for [&Tensor; N] {
292    fn einsum_into(
293        &self,
294        subscripts: &str,
295        session: &mut dyn BackendSession,
296        out: TensorWrite<'_>,
297    ) -> Result<()> {
298        self.as_slice().einsum_into(subscripts, session, out)
299    }
300
301    fn einsum_into_notation(
302        &self,
303        notation: &EinsumNotation,
304        session: &mut dyn BackendSession,
305        out: TensorWrite<'_>,
306    ) -> Result<()> {
307        self.as_slice().einsum_into_notation(notation, session, out)
308    }
309
310    fn einsum_into_subscripts(
311        &self,
312        subscripts: &EinsumSubscripts,
313        session: &mut dyn BackendSession,
314        out: TensorWrite<'_>,
315    ) -> Result<()> {
316        self.as_slice()
317            .einsum_into_subscripts(subscripts, session, out)
318    }
319}
320
321/// Backend-explicit einsum methods for typed concrete tensors.
322///
323/// The result keeps the same scalar type as the inputs. Mixed dtypes should use
324/// [`TensorEinsumExt`] on dtype-erased [`Tensor`] values instead.
325///
326/// # Examples
327///
328/// ```
329/// use tenferro_cpu::CpuBackend;
330/// use tenferro_einsum::TypedTensorEinsumExt;
331/// use tenferro_tensor::{BackendSessionHost, TypedTensor};
332///
333/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
334/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![3, 4], vec![1.0; 12]).unwrap();
335/// let mut backend = CpuBackend::new();
336///
337/// let out = backend.with_backend_session(|session| {
338///     [&lhs, &rhs].einsum("ij,jk->ik", session)
339/// })?;
340/// assert_eq!(out.shape(), &[2, 4]);
341/// # Ok::<(), tenferro_einsum::Error>(())
342/// ```
343pub trait TypedTensorEinsumExt<T: TensorScalar> {
344    /// Execute an einsum from string notation.
345    ///
346    /// # Errors
347    ///
348    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
349    /// [`Error::Validation`] with shape or rank payloads for an invalid
350    /// contraction, or [`Error::Tensor`] for a typed backend failure.
351    fn einsum(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<TypedTensor<T>>;
352
353    /// Execute an einsum from rank-unresolved notation.
354    ///
355    /// # Examples
356    ///
357    /// ```
358    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
359    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
360    /// assert_eq!(notation.input_count(), 1);
361    /// ```
362    ///
363    /// # Errors
364    ///
365    /// Returns a typed validation or backend error when notation, shapes, or execution are invalid.
366    fn einsum_notation(
367        &self,
368        notation: &EinsumNotation,
369        session: &mut dyn BackendSession,
370    ) -> Result<TypedTensor<T>>;
371
372    /// Execute an einsum from parsed integer-label subscripts.
373    ///
374    /// # Errors
375    ///
376    /// Returns [`Error::Validation`] with shape or rank payloads for an invalid
377    /// contraction, or [`Error::Tensor`] for a typed backend failure.
378    fn einsum_subscripts(
379        &self,
380        subscripts: &EinsumSubscripts,
381        session: &mut dyn BackendSession,
382    ) -> Result<TypedTensor<T>>;
383}
384
385impl<T: TensorScalar> TypedTensorEinsumExt<T> for [&TypedTensor<T>] {
386    fn einsum(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<TypedTensor<T>> {
387        let notation = parse_einsum_notation(subscripts)?;
388        self.einsum_notation(&notation, session)
389    }
390
391    fn einsum_notation(
392        &self,
393        notation: &EinsumNotation,
394        session: &mut dyn BackendSession,
395    ) -> Result<TypedTensor<T>> {
396        let subscripts = resolve_typed_notation(self, notation)?;
397        typed_einsum_subscripts(session, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
398    }
399
400    fn einsum_subscripts(
401        &self,
402        subscripts: &EinsumSubscripts,
403        session: &mut dyn BackendSession,
404    ) -> Result<TypedTensor<T>> {
405        let subscripts = Subscripts::from(subscripts);
406        typed_einsum_subscripts(session, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
407    }
408}
409
410impl<T: TensorScalar, const N: usize> TypedTensorEinsumExt<T> for [&TypedTensor<T>; N] {
411    fn einsum(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<TypedTensor<T>> {
412        self.as_slice().einsum(subscripts, session)
413    }
414
415    fn einsum_notation(
416        &self,
417        notation: &EinsumNotation,
418        session: &mut dyn BackendSession,
419    ) -> Result<TypedTensor<T>> {
420        self.as_slice().einsum_notation(notation, session)
421    }
422
423    fn einsum_subscripts(
424        &self,
425        subscripts: &EinsumSubscripts,
426        session: &mut dyn BackendSession,
427    ) -> Result<TypedTensor<T>> {
428        self.as_slice().einsum_subscripts(subscripts, session)
429    }
430}
431
432/// Backend-explicit einsum methods for typed borrowed views.
433///
434/// The `_read` suffix distinguishes this borrowed-view surface from
435/// [`TypedTensorEinsumExt`], whose unsuffixed methods accept only owned compact
436/// typed tensors.
437///
438/// # Examples
439///
440/// ```
441/// use tenferro_cpu::CpuBackend;
442/// use tenferro_einsum::TypedTensorReadEinsumExt;
443/// use tenferro_tensor::{BackendSessionHost, TypedTensor};
444///
445/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 2.0])?;
446/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![3.0, 4.0])?;
447/// let mut backend = CpuBackend::new();
448/// let result = backend.with_backend_session(|session| {
449///     [lhs.as_view(), rhs.as_view()].einsum_read("i,i->", session)
450/// })?;
451/// assert_eq!(result.as_slice()?, &[11.0]);
452/// # Ok::<(), tenferro_einsum::Error>(())
453/// ```
454pub trait TypedTensorReadEinsumExt<T: TensorScalar> {
455    /// Execute an einsum from string notation over typed borrowed views.
456    ///
457    /// # Examples
458    ///
459    /// ```
460    /// use tenferro_cpu::CpuBackend;
461    /// use tenferro_einsum::TypedTensorReadEinsumExt;
462    /// use tenferro_tensor::{BackendSessionHost, TypedTensor};
463    ///
464    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
465    /// let mut backend = CpuBackend::new();
466    /// let result = backend.with_backend_session(|session| {
467    ///     [input.as_view()].einsum_read("i->i", session)
468    /// })?;
469    /// assert_eq!(result.as_slice()?, &[2.0, 3.0]);
470    /// # Ok::<(), tenferro_einsum::Error>(())
471    /// ```
472    ///
473    /// # Errors
474    ///
475    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
476    /// [`Error::Validation`] with shape or rank payloads for incompatible
477    /// views, or [`Error::Tensor`] for a typed backend failure.
478    fn einsum_read(
479        &self,
480        subscripts: &str,
481        session: &mut dyn BackendSession,
482    ) -> Result<TypedTensor<T>>;
483
484    /// Execute an einsum from rank-unresolved notation over typed views.
485    ///
486    /// # Examples
487    ///
488    /// ```
489    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
490    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
491    /// assert_eq!(notation.input_count(), 1);
492    /// ```
493    ///
494    /// # Errors
495    ///
496    /// Returns a typed validation or backend error when notation, views, or execution are invalid.
497    fn einsum_read_notation(
498        &self,
499        notation: &EinsumNotation,
500        session: &mut dyn BackendSession,
501    ) -> Result<TypedTensor<T>>;
502
503    /// Execute an einsum from parsed integer-label subscripts over typed
504    /// borrowed views.
505    ///
506    /// # Examples
507    ///
508    /// ```
509    /// use tenferro_cpu::CpuBackend;
510    /// use tenferro_einsum::{EinsumSubscripts, TypedTensorReadEinsumExt};
511    /// use tenferro_tensor::{BackendSessionHost, TypedTensor};
512    ///
513    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
514    /// let subscripts = EinsumSubscripts::new(&[&[0]], &[0]);
515    /// let mut backend = CpuBackend::new();
516    /// let result = backend.with_backend_session(|session| {
517    ///     [input.as_view()].einsum_read_subscripts(&subscripts, session)
518    /// })?;
519    /// assert_eq!(result.as_slice()?, &[2.0, 3.0]);
520    /// # Ok::<(), tenferro_einsum::Error>(())
521    /// ```
522    ///
523    /// # Errors
524    ///
525    /// Returns [`Error::Validation`] with shape or rank payloads for
526    /// incompatible views, or [`Error::Tensor`] for a typed backend failure.
527    fn einsum_read_subscripts(
528        &self,
529        subscripts: &EinsumSubscripts,
530        session: &mut dyn BackendSession,
531    ) -> Result<TypedTensor<T>>;
532}
533
534impl<'a, T: TensorScalar> TypedTensorReadEinsumExt<T> for [TypedTensorView<'a, T>] {
535    fn einsum_read(
536        &self,
537        subscripts: &str,
538        session: &mut dyn BackendSession,
539    ) -> Result<TypedTensor<T>> {
540        let notation = parse_einsum_notation(subscripts)?;
541        self.einsum_read_notation(&notation, session)
542    }
543
544    fn einsum_read_notation(
545        &self,
546        notation: &EinsumNotation,
547        session: &mut dyn BackendSession,
548    ) -> Result<TypedTensor<T>> {
549        let subscripts = resolve_view_notation(self, notation)?;
550        typed_view_einsum_subscripts(session, self, &subscripts, TYPED_TENSOR_READ_EINSUM_OP)
551    }
552
553    fn einsum_read_subscripts(
554        &self,
555        subscripts: &EinsumSubscripts,
556        session: &mut dyn BackendSession,
557    ) -> Result<TypedTensor<T>> {
558        let subscripts = Subscripts::from(subscripts);
559        typed_view_einsum_subscripts(session, self, &subscripts, TYPED_TENSOR_READ_EINSUM_OP)
560    }
561}
562
563impl<'a, T: TensorScalar, const N: usize> TypedTensorReadEinsumExt<T>
564    for [TypedTensorView<'a, T>; N]
565{
566    fn einsum_read(
567        &self,
568        subscripts: &str,
569        session: &mut dyn BackendSession,
570    ) -> Result<TypedTensor<T>> {
571        self.as_slice().einsum_read(subscripts, session)
572    }
573
574    fn einsum_read_notation(
575        &self,
576        notation: &EinsumNotation,
577        session: &mut dyn BackendSession,
578    ) -> Result<TypedTensor<T>> {
579        self.as_slice().einsum_read_notation(notation, session)
580    }
581
582    fn einsum_read_subscripts(
583        &self,
584        subscripts: &EinsumSubscripts,
585        session: &mut dyn BackendSession,
586    ) -> Result<TypedTensor<T>> {
587        self.as_slice().einsum_read_subscripts(subscripts, session)
588    }
589}
590
591/// Backend-explicit preallocated-output einsum methods for typed concrete tensors.
592pub trait TypedTensorEinsumIntoExt<T: TensorScalar> {
593    /// Execute an einsum from string notation into caller-provided typed output.
594    ///
595    /// # Errors
596    ///
597    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
598    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
599    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
600    fn einsum_into<'out, O>(
601        &self,
602        subscripts: &str,
603        session: &mut dyn BackendSession,
604        out: O,
605    ) -> Result<()>
606    where
607        O: Into<TypedTensorWrite<'out, T>>;
608
609    /// Execute an einsum from rank-unresolved notation into typed output.
610    ///
611    /// # Examples
612    ///
613    /// ```
614    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
615    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
616    /// assert_eq!(notation.input_count(), 1);
617    /// ```
618    ///
619    /// # Errors
620    ///
621    /// Returns a typed validation or backend error when notation, shapes, or output are invalid.
622    fn einsum_into_notation<'out, O>(
623        &self,
624        notation: &EinsumNotation,
625        session: &mut dyn BackendSession,
626        out: O,
627    ) -> Result<()>
628    where
629        O: Into<TypedTensorWrite<'out, T>>;
630
631    /// Execute an einsum from parsed integer-label subscripts into caller-provided typed output.
632    ///
633    /// # Errors
634    ///
635    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
636    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
637    /// failure.
638    fn einsum_into_subscripts<'out, O>(
639        &self,
640        subscripts: &EinsumSubscripts,
641        session: &mut dyn BackendSession,
642        out: O,
643    ) -> Result<()>
644    where
645        O: Into<TypedTensorWrite<'out, T>>;
646}
647
648impl<T: TensorScalar> TypedTensorEinsumIntoExt<T> for [&TypedTensor<T>] {
649    fn einsum_into<'out, O>(
650        &self,
651        subscripts: &str,
652        session: &mut dyn BackendSession,
653        out: O,
654    ) -> Result<()>
655    where
656        O: Into<TypedTensorWrite<'out, T>>,
657    {
658        let notation = parse_einsum_notation(subscripts)?;
659        self.einsum_into_notation(&notation, session, out)
660    }
661
662    fn einsum_into_notation<'out, O>(
663        &self,
664        notation: &EinsumNotation,
665        session: &mut dyn BackendSession,
666        out: O,
667    ) -> Result<()>
668    where
669        O: Into<TypedTensorWrite<'out, T>>,
670    {
671        let subscripts = resolve_typed_notation(self, notation)?;
672        typed_einsum_into_subscripts(
673            session,
674            self,
675            &subscripts,
676            out.into(),
677            TYPED_TENSOR_EINSUM_INTO_OP,
678        )
679    }
680
681    fn einsum_into_subscripts<'out, O>(
682        &self,
683        subscripts: &EinsumSubscripts,
684        session: &mut dyn BackendSession,
685        out: O,
686    ) -> Result<()>
687    where
688        O: Into<TypedTensorWrite<'out, T>>,
689    {
690        let subscripts = Subscripts::from(subscripts);
691        typed_einsum_into_subscripts(
692            session,
693            self,
694            &subscripts,
695            out.into(),
696            TYPED_TENSOR_EINSUM_INTO_OP,
697        )
698    }
699}
700
701impl<T: TensorScalar, const N: usize> TypedTensorEinsumIntoExt<T> for [&TypedTensor<T>; N] {
702    fn einsum_into<'out, O>(
703        &self,
704        subscripts: &str,
705        session: &mut dyn BackendSession,
706        out: O,
707    ) -> Result<()>
708    where
709        O: Into<TypedTensorWrite<'out, T>>,
710    {
711        self.as_slice().einsum_into(subscripts, session, out)
712    }
713
714    fn einsum_into_notation<'out, O>(
715        &self,
716        notation: &EinsumNotation,
717        session: &mut dyn BackendSession,
718        out: O,
719    ) -> Result<()>
720    where
721        O: Into<TypedTensorWrite<'out, T>>,
722    {
723        self.as_slice().einsum_into_notation(notation, session, out)
724    }
725
726    fn einsum_into_subscripts<'out, O>(
727        &self,
728        subscripts: &EinsumSubscripts,
729        session: &mut dyn BackendSession,
730        out: O,
731    ) -> Result<()>
732    where
733        O: Into<TypedTensorWrite<'out, T>>,
734    {
735        self.as_slice()
736            .einsum_into_subscripts(subscripts, session, out)
737    }
738}
739
740/// Backend-explicit preallocated-output einsum methods for typed borrowed views.
741///
742/// # Examples
743///
744/// ```
745/// use tenferro_cpu::CpuBackend;
746/// use tenferro_einsum::TypedTensorReadEinsumIntoExt;
747/// use tenferro_tensor::{BackendSessionHost, TypedTensor};
748///
749/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 2.0])?;
750/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![3.0, 4.0])?;
751/// let mut output = TypedTensor::<f64>::from_vec_col_major(vec![], vec![0.0])?;
752/// let mut backend = CpuBackend::new();
753/// backend.with_backend_session(|session| {
754///     [lhs.as_view(), rhs.as_view()].einsum_read_into("i,i->", session, &mut output)
755/// })?;
756/// assert_eq!(output.as_slice()?, &[11.0]);
757/// # Ok::<(), tenferro_einsum::Error>(())
758/// ```
759pub trait TypedTensorReadEinsumIntoExt<T: TensorScalar> {
760    /// Execute an einsum from string notation over typed borrowed views into a
761    /// caller-provided typed output.
762    ///
763    /// # Examples
764    ///
765    /// ```
766    /// use tenferro_cpu::CpuBackend;
767    /// use tenferro_einsum::TypedTensorReadEinsumIntoExt;
768    /// use tenferro_tensor::{BackendSessionHost, TypedTensor};
769    ///
770    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
771    /// let mut output = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![0.0; 2])?;
772    /// let mut backend = CpuBackend::new();
773    /// backend.with_backend_session(|session| {
774    ///     [input.as_view()].einsum_read_into("i->i", session, &mut output)
775    /// })?;
776    /// assert_eq!(output.as_slice()?, &[2.0, 3.0]);
777    /// # Ok::<(), tenferro_einsum::Error>(())
778    /// ```
779    ///
780    /// # Errors
781    ///
782    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
783    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
784    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
785    fn einsum_read_into<'out, O>(
786        &self,
787        subscripts: &str,
788        session: &mut dyn BackendSession,
789        out: O,
790    ) -> Result<()>
791    where
792        O: Into<TypedTensorWrite<'out, T>>;
793
794    /// Execute an einsum from rank-unresolved notation over typed views into output.
795    ///
796    /// # Examples
797    ///
798    /// ```
799    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
800    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
801    /// assert_eq!(notation.input_count(), 1);
802    /// ```
803    ///
804    /// # Errors
805    ///
806    /// Returns a typed validation or backend error when notation, views, or output are invalid.
807    fn einsum_read_into_notation<'out, O>(
808        &self,
809        notation: &EinsumNotation,
810        session: &mut dyn BackendSession,
811        out: O,
812    ) -> Result<()>
813    where
814        O: Into<TypedTensorWrite<'out, T>>;
815
816    /// Execute an einsum from parsed integer-label subscripts over typed
817    /// borrowed views into a caller-provided typed output.
818    ///
819    /// # Examples
820    ///
821    /// ```
822    /// use tenferro_cpu::CpuBackend;
823    /// use tenferro_einsum::{EinsumSubscripts, TypedTensorReadEinsumIntoExt};
824    /// use tenferro_tensor::{BackendSessionHost, TypedTensor};
825    ///
826    /// let input = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 3.0])?;
827    /// let mut output = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![0.0; 2])?;
828    /// let subscripts = EinsumSubscripts::new(&[&[0]], &[0]);
829    /// let mut backend = CpuBackend::new();
830    /// backend.with_backend_session(|session| {
831    ///     [input.as_view()].einsum_read_into_subscripts(&subscripts, session, &mut output)
832    /// })?;
833    /// assert_eq!(output.as_slice()?, &[2.0, 3.0]);
834    /// # Ok::<(), tenferro_einsum::Error>(())
835    /// ```
836    ///
837    /// # Errors
838    ///
839    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
840    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
841    /// failure.
842    fn einsum_read_into_subscripts<'out, O>(
843        &self,
844        subscripts: &EinsumSubscripts,
845        session: &mut dyn BackendSession,
846        out: O,
847    ) -> Result<()>
848    where
849        O: Into<TypedTensorWrite<'out, T>>;
850}
851
852impl<'a, T: TensorScalar> TypedTensorReadEinsumIntoExt<T> for [TypedTensorView<'a, T>] {
853    fn einsum_read_into<'out, O>(
854        &self,
855        subscripts: &str,
856        session: &mut dyn BackendSession,
857        out: O,
858    ) -> Result<()>
859    where
860        O: Into<TypedTensorWrite<'out, T>>,
861    {
862        let notation = parse_einsum_notation(subscripts)?;
863        self.einsum_read_into_notation(&notation, session, out)
864    }
865
866    fn einsum_read_into_notation<'out, O>(
867        &self,
868        notation: &EinsumNotation,
869        session: &mut dyn BackendSession,
870        out: O,
871    ) -> Result<()>
872    where
873        O: Into<TypedTensorWrite<'out, T>>,
874    {
875        let subscripts = resolve_view_notation(self, notation)?;
876        typed_view_einsum_into_subscripts(
877            session,
878            self,
879            &subscripts,
880            out.into(),
881            TYPED_TENSOR_READ_EINSUM_INTO_OP,
882        )
883    }
884
885    fn einsum_read_into_subscripts<'out, O>(
886        &self,
887        subscripts: &EinsumSubscripts,
888        session: &mut dyn BackendSession,
889        out: O,
890    ) -> Result<()>
891    where
892        O: Into<TypedTensorWrite<'out, T>>,
893    {
894        let subscripts = Subscripts::from(subscripts);
895        typed_view_einsum_into_subscripts(
896            session,
897            self,
898            &subscripts,
899            out.into(),
900            TYPED_TENSOR_READ_EINSUM_INTO_OP,
901        )
902    }
903}
904
905impl<'a, T: TensorScalar, const N: usize> TypedTensorReadEinsumIntoExt<T>
906    for [TypedTensorView<'a, T>; N]
907{
908    fn einsum_read_into<'out, O>(
909        &self,
910        subscripts: &str,
911        session: &mut dyn BackendSession,
912        out: O,
913    ) -> Result<()>
914    where
915        O: Into<TypedTensorWrite<'out, T>>,
916    {
917        self.as_slice().einsum_read_into(subscripts, session, out)
918    }
919
920    fn einsum_read_into_notation<'out, O>(
921        &self,
922        notation: &EinsumNotation,
923        session: &mut dyn BackendSession,
924        out: O,
925    ) -> Result<()>
926    where
927        O: Into<TypedTensorWrite<'out, T>>,
928    {
929        self.as_slice()
930            .einsum_read_into_notation(notation, session, out)
931    }
932
933    fn einsum_read_into_subscripts<'out, O>(
934        &self,
935        subscripts: &EinsumSubscripts,
936        session: &mut dyn BackendSession,
937        out: O,
938    ) -> Result<()>
939    where
940        O: Into<TypedTensorWrite<'out, T>>,
941    {
942        self.as_slice()
943            .einsum_read_into_subscripts(subscripts, session, out)
944    }
945}
946
947/// Backend-explicit einsum methods for [`TensorRead`] inputs.
948///
949/// Use this surface when an input is a borrowed tensor view rather than an
950/// owned compact [`Tensor`]. The `_read` suffix follows the repository-wide
951/// convention for APIs that explicitly accept read-oriented borrowed inputs.
952///
953/// # Examples
954///
955/// ```
956/// use tenferro_cpu::CpuBackend;
957/// use tenferro_einsum::TensorReadEinsumExt;
958/// use tenferro_tensor::{BackendSessionHost, Tensor, TensorRead, TensorView};
959///
960/// let shape = [2, 3];
961/// let data = [1.0_f64; 6];
962/// let rhs = Tensor::from_vec_col_major(vec![3], vec![1.0_f64; 3]).unwrap();
963/// let inputs = [
964///     TensorRead::from_view(TensorView::f64(&shape, &data)?),
965///     TensorRead::from_tensor(&rhs),
966/// ];
967/// let mut backend = CpuBackend::new();
968///
969/// let out = backend.with_backend_session(|session| inputs.einsum_read("ij,j->i", session))?;
970/// assert_eq!(out.shape(), &[2]);
971/// # Ok::<(), tenferro_einsum::Error>(())
972/// ```
973pub trait TensorReadEinsumExt {
974    /// Execute an einsum from string notation over read-only tensor inputs.
975    ///
976    /// # Errors
977    ///
978    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
979    /// [`Error::Validation`] with shape, rank, or dtype payloads for an invalid
980    /// contraction, or [`Error::Tensor`] for a typed backend failure.
981    fn einsum_read(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<Tensor>;
982
983    /// Execute an einsum from rank-unresolved notation over read-only inputs.
984    ///
985    /// # Examples
986    ///
987    /// ```
988    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
989    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
990    /// assert_eq!(notation.input_count(), 1);
991    /// ```
992    ///
993    /// # Errors
994    ///
995    /// Returns a typed validation or backend error when notation, shapes, or execution are invalid.
996    fn einsum_read_notation(
997        &self,
998        notation: &EinsumNotation,
999        session: &mut dyn BackendSession,
1000    ) -> Result<Tensor>;
1001
1002    /// Execute an einsum from parsed integer-label subscripts over read-only
1003    /// tensor inputs.
1004    ///
1005    /// # Errors
1006    ///
1007    /// Returns [`Error::Validation`] with shape, rank, or dtype payloads for an
1008    /// invalid contraction, or [`Error::Tensor`] for a typed backend failure.
1009    fn einsum_read_subscripts(
1010        &self,
1011        subscripts: &EinsumSubscripts,
1012        session: &mut dyn BackendSession,
1013    ) -> Result<Tensor>;
1014}
1015
1016impl<'a> TensorReadEinsumExt for [TensorRead<'a>] {
1017    fn einsum_read(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<Tensor> {
1018        let notation = parse_einsum_notation(subscripts)?;
1019        self.einsum_read_notation(&notation, session)
1020    }
1021
1022    fn einsum_read_notation(
1023        &self,
1024        notation: &EinsumNotation,
1025        session: &mut dyn BackendSession,
1026    ) -> Result<Tensor> {
1027        let subscripts = resolve_read_notation(self, notation)?;
1028        eager_einsum_read_subscripts_on_session(session, self, &subscripts).map_err(Error::from)
1029    }
1030
1031    fn einsum_read_subscripts(
1032        &self,
1033        subscripts: &EinsumSubscripts,
1034        session: &mut dyn BackendSession,
1035    ) -> Result<Tensor> {
1036        let subscripts = Subscripts::from(subscripts);
1037        eager_einsum_read_subscripts_on_session(session, self, &subscripts).map_err(Error::from)
1038    }
1039}
1040
1041impl<'a, const N: usize> TensorReadEinsumExt for [TensorRead<'a>; N] {
1042    fn einsum_read(&self, subscripts: &str, session: &mut dyn BackendSession) -> Result<Tensor> {
1043        self.as_slice().einsum_read(subscripts, session)
1044    }
1045
1046    fn einsum_read_notation(
1047        &self,
1048        notation: &EinsumNotation,
1049        session: &mut dyn BackendSession,
1050    ) -> Result<Tensor> {
1051        self.as_slice().einsum_read_notation(notation, session)
1052    }
1053
1054    fn einsum_read_subscripts(
1055        &self,
1056        subscripts: &EinsumSubscripts,
1057        session: &mut dyn BackendSession,
1058    ) -> Result<Tensor> {
1059        self.as_slice().einsum_read_subscripts(subscripts, session)
1060    }
1061}
1062
1063/// Backend-explicit preallocated-output einsum methods for [`TensorRead`] inputs.
1064pub trait TensorReadEinsumIntoExt {
1065    /// Execute an einsum from string notation over read-only inputs into caller-provided output.
1066    ///
1067    /// # Errors
1068    ///
1069    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
1070    /// [`Error::Validation`] with a shape, rank, or dtype payload when inputs or
1071    /// output do not match, or [`Error::Tensor`] for a typed backend failure.
1072    fn einsum_read_into(
1073        &self,
1074        subscripts: &str,
1075        session: &mut dyn BackendSession,
1076        out: TensorWrite<'_>,
1077    ) -> Result<()>;
1078
1079    /// Execute an einsum from rank-unresolved notation over read-only inputs into output.
1080    ///
1081    /// # Examples
1082    ///
1083    /// ```
1084    /// use tenferro_einsum::{EinsumAxis, EinsumNotation};
1085    /// let notation = EinsumNotation::new(&[&[EinsumAxis::Ellipsis]], &[]);
1086    /// assert_eq!(notation.input_count(), 1);
1087    /// ```
1088    ///
1089    /// # Errors
1090    ///
1091    /// Returns a typed validation or backend error when notation, shapes, or output are invalid.
1092    fn einsum_read_into_notation(
1093        &self,
1094        notation: &EinsumNotation,
1095        session: &mut dyn BackendSession,
1096        out: TensorWrite<'_>,
1097    ) -> Result<()>;
1098
1099    /// Execute an einsum from parsed integer-label subscripts over read-only inputs into output.
1100    ///
1101    /// # Errors
1102    ///
1103    /// Returns [`Error::Validation`] with a shape, rank, or dtype payload when
1104    /// inputs or output do not match, or [`Error::Tensor`] for a typed backend
1105    /// failure.
1106    fn einsum_read_into_subscripts(
1107        &self,
1108        subscripts: &EinsumSubscripts,
1109        session: &mut dyn BackendSession,
1110        out: TensorWrite<'_>,
1111    ) -> Result<()>;
1112}
1113
1114impl<'a> TensorReadEinsumIntoExt for [TensorRead<'a>] {
1115    fn einsum_read_into(
1116        &self,
1117        subscripts: &str,
1118        session: &mut dyn BackendSession,
1119        out: TensorWrite<'_>,
1120    ) -> Result<()> {
1121        let notation = parse_einsum_notation(subscripts)?;
1122        self.einsum_read_into_notation(&notation, session, out)
1123    }
1124
1125    fn einsum_read_into_notation(
1126        &self,
1127        notation: &EinsumNotation,
1128        session: &mut dyn BackendSession,
1129        out: TensorWrite<'_>,
1130    ) -> Result<()> {
1131        let subscripts = resolve_read_notation(self, notation)?;
1132        tensor_read_einsum_into_subscripts(
1133            session,
1134            self,
1135            &subscripts,
1136            out,
1137            TENSOR_READ_EINSUM_INTO_OP,
1138        )
1139    }
1140
1141    fn einsum_read_into_subscripts(
1142        &self,
1143        subscripts: &EinsumSubscripts,
1144        session: &mut dyn BackendSession,
1145        out: TensorWrite<'_>,
1146    ) -> Result<()> {
1147        let subscripts = Subscripts::from(subscripts);
1148        tensor_read_einsum_into_subscripts(
1149            session,
1150            self,
1151            &subscripts,
1152            out,
1153            TENSOR_READ_EINSUM_INTO_OP,
1154        )
1155    }
1156}
1157
1158impl<'a, const N: usize> TensorReadEinsumIntoExt for [TensorRead<'a>; N] {
1159    fn einsum_read_into(
1160        &self,
1161        subscripts: &str,
1162        session: &mut dyn BackendSession,
1163        out: TensorWrite<'_>,
1164    ) -> Result<()> {
1165        self.as_slice().einsum_read_into(subscripts, session, out)
1166    }
1167
1168    fn einsum_read_into_notation(
1169        &self,
1170        notation: &EinsumNotation,
1171        session: &mut dyn BackendSession,
1172        out: TensorWrite<'_>,
1173    ) -> Result<()> {
1174        self.as_slice()
1175            .einsum_read_into_notation(notation, session, out)
1176    }
1177
1178    fn einsum_read_into_subscripts(
1179        &self,
1180        subscripts: &EinsumSubscripts,
1181        session: &mut dyn BackendSession,
1182        out: TensorWrite<'_>,
1183    ) -> Result<()> {
1184        self.as_slice()
1185            .einsum_read_into_subscripts(subscripts, session, out)
1186    }
1187}
1188
1189/// Prepared concrete einsum plan for repeated executions with fixed input
1190/// dtype and shape metadata.
1191///
1192/// Preparing a plan parses and optimizes the contraction tree once. Execution
1193/// validates the later inputs against the prepared dtype and shape contract,
1194/// then runs the stored tree without re-planning.
1195///
1196/// # Examples
1197///
1198/// ```
1199/// use tenferro_cpu::CpuBackend;
1200/// use tenferro_einsum::ConcreteEinsumPlan;
1201/// use tenferro_tensor::{BackendSessionHost, Tensor};
1202///
1203/// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
1204/// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
1205/// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik")?;
1206///
1207/// let mut backend = CpuBackend::new();
1208/// let out = backend
1209///     .with_backend_session(|session| plan.execute([&lhs, &rhs], session))?;
1210/// assert_eq!(out.shape(), &[2, 4]);
1211/// # Ok::<(), tenferro_einsum::Error>(())
1212/// ```
1213#[derive(Debug)]
1214pub struct ConcreteEinsumPlan {
1215    tree: ContractionTree,
1216    inputs: Vec<ConcreteEinsumInputSpec>,
1217}
1218
1219impl ConcreteEinsumPlan {
1220    /// Prepare a plan from dtype-erased concrete tensor inputs and string
1221    /// notation.
1222    ///
1223    /// # Errors
1224    ///
1225    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
1226    /// [`Error::Validation`] for rank, shape, or dtype contract violations, or
1227    /// [`Error::Planning`] when no valid contraction tree can be built.
1228    pub fn prepare<'a, I>(inputs: I, subscripts: &str) -> Result<Self>
1229    where
1230        I: AsRef<[&'a Tensor]>,
1231    {
1232        let notation = parse_einsum_notation(subscripts)?;
1233        Self::prepare_notation(inputs, &notation)
1234    }
1235
1236    /// Prepare a plan from dtype-erased concrete tensor inputs and parsed
1237    /// integer-label subscripts.
1238    ///
1239    /// # Errors
1240    ///
1241    /// Returns [`Error::Validation`] for rank, shape, or dtype contract
1242    /// violations, or [`Error::Planning`] when no valid contraction tree can be
1243    /// built.
1244    pub fn prepare_subscripts<'a, I>(inputs: I, subscripts: &EinsumSubscripts) -> Result<Self>
1245    where
1246        I: AsRef<[&'a Tensor]>,
1247    {
1248        let subscripts = Subscripts::from(subscripts);
1249        Self::prepare_subscripts_internal(input_specs(inputs.as_ref()), &subscripts)
1250    }
1251
1252    /// Prepare a plan from rank-unresolved notation and concrete tensor inputs.
1253    ///
1254    /// # Errors
1255    ///
1256    /// Returns [`Error::InvalidSubscripts`] for malformed axis tokens,
1257    /// [`Error::Validation`] for rank, shape, or dtype violations, or
1258    /// [`Error::Planning`] when no contraction tree can be built.
1259    pub fn prepare_notation<'a, I>(inputs: I, notation: &EinsumNotation) -> Result<Self>
1260    where
1261        I: AsRef<[&'a Tensor]>,
1262    {
1263        let inputs = inputs.as_ref();
1264        let subscripts = resolve_tensor_notation(inputs, notation)?;
1265        Self::prepare_subscripts_internal(input_specs(inputs), &subscripts)
1266    }
1267
1268    /// Prepare a plan from typed concrete tensor inputs and string notation.
1269    ///
1270    /// # Errors
1271    ///
1272    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
1273    /// [`Error::Validation`] for rank or shape contract violations, or
1274    /// [`Error::Planning`] when no valid contraction tree can be built.
1275    pub fn prepare_typed<'a, T, I>(inputs: I, subscripts: &str) -> Result<Self>
1276    where
1277        T: TensorScalar,
1278        I: AsRef<[&'a TypedTensor<T>]>,
1279    {
1280        let notation = parse_einsum_notation(subscripts)?;
1281        Self::prepare_typed_notation(inputs, &notation)
1282    }
1283
1284    /// Prepare a plan from typed concrete tensor inputs and parsed integer-label
1285    /// subscripts.
1286    ///
1287    /// # Errors
1288    ///
1289    /// Returns [`Error::Validation`] for rank or shape contract violations, or
1290    /// [`Error::Planning`] when no valid contraction tree can be built.
1291    pub fn prepare_typed_subscripts<'a, T, I>(
1292        inputs: I,
1293        subscripts: &EinsumSubscripts,
1294    ) -> Result<Self>
1295    where
1296        T: TensorScalar,
1297        I: AsRef<[&'a TypedTensor<T>]>,
1298    {
1299        let subscripts = Subscripts::from(subscripts);
1300        Self::prepare_subscripts_internal(typed_input_specs(inputs.as_ref()), &subscripts)
1301    }
1302
1303    /// Prepare a plan from rank-unresolved notation and typed concrete inputs.
1304    ///
1305    /// # Errors
1306    ///
1307    /// Returns [`Error::InvalidSubscripts`] for malformed axis tokens,
1308    /// [`Error::Validation`] for rank or shape violations, or [`Error::Planning`]
1309    /// when no contraction tree can be built.
1310    pub fn prepare_typed_notation<'a, T, I>(inputs: I, notation: &EinsumNotation) -> Result<Self>
1311    where
1312        T: TensorScalar,
1313        I: AsRef<[&'a TypedTensor<T>]>,
1314    {
1315        let inputs = inputs.as_ref();
1316        let subscripts = resolve_typed_notation(inputs, notation)?;
1317        Self::prepare_subscripts_internal(typed_input_specs(inputs), &subscripts)
1318    }
1319
1320    /// Prepare a plan from read-only tensor inputs and string notation.
1321    ///
1322    /// # Errors
1323    ///
1324    /// Returns [`Error::InvalidSubscripts`] for malformed notation,
1325    /// [`Error::Validation`] for rank, shape, or dtype contract violations, or
1326    /// [`Error::Planning`] when no valid contraction tree can be built.
1327    pub fn prepare_read<'a, I>(inputs: I, subscripts: &str) -> Result<Self>
1328    where
1329        I: AsRef<[TensorRead<'a>]>,
1330    {
1331        let notation = parse_einsum_notation(subscripts)?;
1332        Self::prepare_read_notation(inputs, &notation)
1333    }
1334
1335    /// Prepare a plan from read-only tensor inputs and parsed integer-label
1336    /// subscripts.
1337    ///
1338    /// # Errors
1339    ///
1340    /// Returns [`Error::Validation`] for rank, shape, or dtype contract
1341    /// violations, or [`Error::Planning`] when no valid contraction tree can be
1342    /// built.
1343    pub fn prepare_read_subscripts<'a, I>(inputs: I, subscripts: &EinsumSubscripts) -> Result<Self>
1344    where
1345        I: AsRef<[TensorRead<'a>]>,
1346    {
1347        let subscripts = Subscripts::from(subscripts);
1348        Self::prepare_subscripts_internal(read_input_specs(inputs.as_ref()), &subscripts)
1349    }
1350
1351    /// Prepare a plan from rank-unresolved notation and read-only inputs.
1352    ///
1353    /// # Errors
1354    ///
1355    /// Returns [`Error::InvalidSubscripts`] for malformed axis tokens,
1356    /// [`Error::Validation`] for rank, shape, or dtype violations, or
1357    /// [`Error::Planning`] when no contraction tree can be built.
1358    pub fn prepare_read_notation<'a, I>(inputs: I, notation: &EinsumNotation) -> Result<Self>
1359    where
1360        I: AsRef<[TensorRead<'a>]>,
1361    {
1362        let inputs = inputs.as_ref();
1363        let subscripts = resolve_read_notation(inputs, notation)?;
1364        Self::prepare_subscripts_internal(read_input_specs(inputs), &subscripts)
1365    }
1366
1367    /// Number of binary contraction steps in the prepared tree (diagnostics).
1368    pub(crate) fn step_count(&self) -> usize {
1369        self.tree.step_count()
1370    }
1371
1372    /// Execute this plan on dtype-erased concrete tensor inputs inside a
1373    /// borrowed backend session.
1374    ///
1375    /// Validation and the contraction itself run in the caller's `session`;
1376    /// this method never enters a new backend session.
1377    ///
1378    /// # Examples
1379    ///
1380    /// ```
1381    /// use tenferro_cpu::CpuBackend;
1382    /// use tenferro_einsum::ConcreteEinsumPlan;
1383    /// use tenferro_tensor::{BackendSessionHost, Tensor};
1384    ///
1385    /// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
1386    /// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
1387    /// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik")?;
1388    ///
1389    /// let mut backend = CpuBackend::new();
1390    /// let out = backend
1391    ///     .with_backend_session(|session| plan.execute([&lhs, &rhs], session))?;
1392    /// assert_eq!(out.shape(), &[2, 4]);
1393    /// # Ok::<(), tenferro_einsum::Error>(())
1394    /// ```
1395    ///
1396    /// # Errors
1397    ///
1398    /// Returns [`Error::Validation`] when inputs violate the prepared rank,
1399    /// shape, or input-count contract, [`Error::Tensor`] with a
1400    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload when an
1401    /// input dtype differs from the prepared contract, or [`Error::Tensor`]
1402    /// for a typed backend failure.
1403    pub fn execute<'a, I>(&self, inputs: I, session: &mut dyn BackendSession) -> Result<Tensor>
1404    where
1405        I: AsRef<[&'a Tensor]>,
1406    {
1407        let inputs = inputs.as_ref();
1408        self.validate_inputs(&input_specs(inputs), PLAN_EXECUTE_OP)?;
1409        eager_einsum_exec(session, inputs, &self.tree).map_err(Error::from)
1410    }
1411
1412    /// Execute this plan on typed concrete tensor inputs inside a borrowed
1413    /// backend session.
1414    ///
1415    /// Validation and the contraction itself run in the caller's `session`;
1416    /// this method never enters a new backend session.
1417    ///
1418    /// # Examples
1419    ///
1420    /// ```
1421    /// use tenferro_cpu::CpuBackend;
1422    /// use tenferro_einsum::ConcreteEinsumPlan;
1423    /// use tenferro_tensor::{BackendSessionHost, TypedTensor};
1424    ///
1425    /// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
1426    /// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![3, 4], vec![1.0; 12]).unwrap();
1427    /// let plan = ConcreteEinsumPlan::prepare_typed([&lhs, &rhs], "ij,jk->ik")?;
1428    ///
1429    /// let mut backend = CpuBackend::new();
1430    /// let out = backend
1431    ///     .with_backend_session(|session| plan.execute_typed([&lhs, &rhs], session))?;
1432    /// assert_eq!(out.shape(), &[2, 4]);
1433    /// # Ok::<(), tenferro_einsum::Error>(())
1434    /// ```
1435    ///
1436    /// # Errors
1437    ///
1438    /// Returns [`Error::Validation`] when inputs violate the prepared rank,
1439    /// shape, or input-count contract, [`Error::Tensor`] with a
1440    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload when the
1441    /// prepared dtype differs from `T` or the eager result dtype, or
1442    /// [`Error::Tensor`] for a typed backend failure.
1443    pub fn execute_typed<'a, T, I>(
1444        &self,
1445        inputs: I,
1446        session: &mut dyn BackendSession,
1447    ) -> Result<TypedTensor<T>>
1448    where
1449        T: TensorScalar,
1450        I: AsRef<[&'a TypedTensor<T>]>,
1451    {
1452        let inputs = inputs.as_ref();
1453        self.validate_inputs(&typed_input_specs(inputs), PLAN_EXECUTE_OP)?;
1454        let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1455        let result = eager_einsum_exec_read(session, &reads, &self.tree)?;
1456        into_typed_result(result, PLAN_EXECUTE_OP)
1457    }
1458
1459    /// Execute this plan on read-only tensor inputs inside a borrowed backend
1460    /// session.
1461    ///
1462    /// Validation and the contraction itself run in the caller's `session`;
1463    /// this method never enters a new backend session.
1464    ///
1465    /// # Examples
1466    ///
1467    /// ```
1468    /// use tenferro_cpu::CpuBackend;
1469    /// use tenferro_einsum::ConcreteEinsumPlan;
1470    /// use tenferro_tensor::{BackendSessionHost, Tensor, TensorRead};
1471    ///
1472    /// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
1473    /// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
1474    /// let plan = ConcreteEinsumPlan::prepare_read(
1475    ///     [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
1476    ///     "ij,jk->ik",
1477    /// )?;
1478    ///
1479    /// let mut backend = CpuBackend::new();
1480    /// let reads = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
1481    /// let out = backend
1482    ///     .with_backend_session(|session| plan.execute_read(reads, session))?;
1483    /// assert_eq!(out.shape(), &[2, 4]);
1484    /// # Ok::<(), tenferro_einsum::Error>(())
1485    /// ```
1486    ///
1487    /// # Errors
1488    ///
1489    /// Returns [`Error::Validation`] when inputs violate the prepared rank,
1490    /// shape, or input-count contract, [`Error::Tensor`] with a
1491    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload when an
1492    /// input dtype differs from the prepared contract, or [`Error::Tensor`]
1493    /// for a typed backend failure.
1494    pub fn execute_read<'a, I>(&self, inputs: I, session: &mut dyn BackendSession) -> Result<Tensor>
1495    where
1496        I: AsRef<[TensorRead<'a>]>,
1497    {
1498        let inputs = inputs.as_ref();
1499        self.validate_inputs(&read_input_specs(inputs), PLAN_EXECUTE_OP)?;
1500        eager_einsum_exec_read(session, inputs, &self.tree).map_err(Error::from)
1501    }
1502
1503    /// Execute this plan on dtype-erased concrete tensor inputs into
1504    /// caller-provided output inside a borrowed backend session.
1505    ///
1506    /// Validation and the contraction itself run in the caller's `session`;
1507    /// this method never enters a new backend session.
1508    ///
1509    /// # Examples
1510    ///
1511    /// ```
1512    /// use tenferro_cpu::CpuBackend;
1513    /// use tenferro_einsum::ConcreteEinsumPlan;
1514    /// use tenferro_tensor::{BackendSessionHost, Tensor, TensorWrite};
1515    ///
1516    /// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
1517    /// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
1518    /// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik")?;
1519    ///
1520    /// let mut backend = CpuBackend::new();
1521    /// let mut out = Tensor::from_vec_col_major(vec![2, 4], vec![0.0_f64; 8]).unwrap();
1522    /// backend.with_backend_session(|session| {
1523    ///     plan.execute_into(
1524    ///         [&lhs, &rhs],
1525    ///         session,
1526    ///         TensorWrite::from_tensor(&mut out),
1527    ///     )
1528    /// })?;
1529    /// assert_eq!(out.as_slice::<f64>()?, vec![3.0_f64; 8].as_slice());
1530    /// # Ok::<(), tenferro_einsum::Error>(())
1531    /// ```
1532    ///
1533    /// # Errors
1534    ///
1535    /// Returns [`Error::Validation`] for input or output rank, shape, or
1536    /// input-count contract violations, [`Error::Tensor`] with a
1537    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload for dtype
1538    /// mismatches, or [`Error::Tensor`] for a typed backend failure.
1539    pub fn execute_into<'a, I>(
1540        &self,
1541        inputs: I,
1542        session: &mut dyn BackendSession,
1543        out: TensorWrite<'_>,
1544    ) -> Result<()>
1545    where
1546        I: AsRef<[&'a Tensor]>,
1547    {
1548        let inputs = inputs.as_ref();
1549        let specs = input_specs(inputs);
1550        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1551        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1552        let reads: Vec<_> = inputs
1553            .iter()
1554            .map(|tensor| TensorRead::from_tensor(tensor))
1555            .collect();
1556        eager_einsum_exec_read_into(session, &reads, &self.tree, out).map_err(Error::from)
1557    }
1558
1559    /// Execute this plan on typed concrete tensor inputs into caller-provided
1560    /// output inside a borrowed backend session.
1561    ///
1562    /// Validation and the contraction itself run in the caller's `session`;
1563    /// this method never enters a new backend session.
1564    ///
1565    /// # Examples
1566    ///
1567    /// ```
1568    /// use tenferro_cpu::CpuBackend;
1569    /// use tenferro_einsum::ConcreteEinsumPlan;
1570    /// use tenferro_tensor::{BackendSessionHost, TypedTensor};
1571    ///
1572    /// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
1573    /// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![3, 4], vec![1.0; 12]).unwrap();
1574    /// let plan = ConcreteEinsumPlan::prepare_typed([&lhs, &rhs], "ij,jk->ik")?;
1575    ///
1576    /// let mut backend = CpuBackend::new();
1577    /// let mut out = TypedTensor::<f64>::from_vec_col_major(vec![2, 4], vec![0.0; 8]).unwrap();
1578    /// backend.with_backend_session(|session| {
1579    ///     plan.execute_typed_into([&lhs, &rhs], session, &mut out)
1580    /// })?;
1581    /// assert_eq!(out.as_slice()?, vec![3.0_f64; 8].as_slice());
1582    /// # Ok::<(), tenferro_einsum::Error>(())
1583    /// ```
1584    ///
1585    /// # Errors
1586    ///
1587    /// Returns [`Error::Validation`] for input or output rank, shape, or
1588    /// input-count contract violations, [`Error::Tensor`] with a
1589    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload when the
1590    /// prepared dtype differs from `T` or the output dtype, or
1591    /// [`Error::Tensor`] for a typed backend failure.
1592    pub fn execute_typed_into<'a, 'out, T, I, O>(
1593        &self,
1594        inputs: I,
1595        session: &mut dyn BackendSession,
1596        out: O,
1597    ) -> Result<()>
1598    where
1599        T: TensorScalar,
1600        I: AsRef<[&'a TypedTensor<T>]>,
1601        O: Into<TypedTensorWrite<'out, T>>,
1602    {
1603        let inputs = inputs.as_ref();
1604        let specs = typed_input_specs(inputs);
1605        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1606        let out = out.into().into_tensor_write();
1607        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1608        let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1609        eager_einsum_exec_read_into(session, &reads, &self.tree, out).map_err(Error::from)
1610    }
1611
1612    /// Execute this plan on read-only tensor inputs into caller-provided output
1613    /// inside a borrowed backend session.
1614    ///
1615    /// Validation and the contraction itself run in the caller's `session`;
1616    /// this method never enters a new backend session.
1617    ///
1618    /// # Examples
1619    ///
1620    /// ```
1621    /// use tenferro_cpu::CpuBackend;
1622    /// use tenferro_einsum::ConcreteEinsumPlan;
1623    /// use tenferro_tensor::{BackendSessionHost, Tensor, TensorRead, TensorWrite};
1624    ///
1625    /// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
1626    /// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
1627    /// let plan = ConcreteEinsumPlan::prepare_read(
1628    ///     [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
1629    ///     "ij,jk->ik",
1630    /// )?;
1631    ///
1632    /// let mut backend = CpuBackend::new();
1633    /// let mut out = Tensor::from_vec_col_major(vec![2, 4], vec![0.0_f64; 8]).unwrap();
1634    /// let reads = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
1635    /// backend.with_backend_session(|session| {
1636    ///     plan.execute_read_into(reads, session, TensorWrite::from_tensor(&mut out))
1637    /// })?;
1638    /// assert_eq!(out.as_slice::<f64>()?, vec![3.0_f64; 8].as_slice());
1639    /// # Ok::<(), tenferro_einsum::Error>(())
1640    /// ```
1641    ///
1642    /// # Errors
1643    ///
1644    /// Returns [`Error::Validation`] for input or output rank, shape, or
1645    /// input-count contract violations, [`Error::Tensor`] with a
1646    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload for dtype
1647    /// mismatches, or [`Error::Tensor`] for a typed backend failure.
1648    pub fn execute_read_into<'a, I>(
1649        &self,
1650        inputs: I,
1651        session: &mut dyn BackendSession,
1652        out: TensorWrite<'_>,
1653    ) -> Result<()>
1654    where
1655        I: AsRef<[TensorRead<'a>]>,
1656    {
1657        let inputs = inputs.as_ref();
1658        let specs = read_input_specs(inputs);
1659        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1660        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1661        eager_einsum_exec_read_into(session, inputs, &self.tree, out).map_err(Error::from)
1662    }
1663
1664    /// Execute this plan on read-only inputs with scaled output accumulation
1665    /// inside a borrowed backend session.
1666    ///
1667    /// `accumulation` follows the dot-general contract:
1668    /// `out = alpha * einsum(inputs) + beta * out`.
1669    ///
1670    /// # Examples
1671    ///
1672    /// ```
1673    /// use tenferro_cpu::CpuBackend;
1674    /// use tenferro_einsum::ConcreteEinsumPlan;
1675    /// use tenferro_tensor::{
1676    ///     BackendSessionHost, DotGeneralAccumulation, DType, Tensor, TensorRead, TensorWrite,
1677    /// };
1678    ///
1679    /// let lhs = Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?;
1680    /// let rhs = Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?;
1681    /// let mut out = Tensor::from_vec_col_major(vec![], vec![1.0_f64])?;
1682    /// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "i,i->")?;
1683    /// let mut backend = CpuBackend::new();
1684    /// backend.with_backend_session(|session| {
1685    ///     plan.execute_read_into_accum(
1686    ///         [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
1687    ///         session,
1688    ///         DotGeneralAccumulation::add_to(DType::F64)?,
1689    ///         TensorWrite::from_tensor(&mut out),
1690    ///     )
1691    /// })?;
1692    /// assert_eq!(out.as_slice::<f64>()?, &[7.0]);
1693    /// # Ok::<(), tenferro_einsum::Error>(())
1694    /// ```
1695    ///
1696    /// # Errors
1697    ///
1698    /// Returns [`Error::Validation`] for input or output rank, shape, or
1699    /// input-count contract violations, [`Error::Tensor`] with a
1700    /// `tenferro_tensor::Error::Validation` `DTypeMismatch` payload for dtype
1701    /// mismatches, [`Error::Numerical`] for an invalid accumulation, or
1702    /// [`Error::Tensor`] for a typed backend failure.
1703    pub fn execute_read_into_accum<'a, I>(
1704        &self,
1705        inputs: I,
1706        session: &mut dyn BackendSession,
1707        accumulation: DotGeneralAccumulation,
1708        out: TensorWrite<'_>,
1709    ) -> Result<()>
1710    where
1711        I: AsRef<[TensorRead<'a>]>,
1712    {
1713        let inputs = inputs.as_ref();
1714        let specs = read_input_specs(inputs);
1715        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
1716        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
1717        eager_einsum_exec_read_into_accum(session, inputs, &self.tree, accumulation, out)
1718            .map_err(Error::from)
1719    }
1720
1721    fn prepare_subscripts_internal(
1722        inputs: Vec<ConcreteEinsumInputSpec>,
1723        subscripts: &Subscripts,
1724    ) -> Result<Self> {
1725        let shapes: Vec<&[usize]> = inputs.iter().map(|input| input.shape.as_slice()).collect();
1726        let tree = plan_subscripts(subscripts, &shapes)?;
1727        Ok(Self { tree, inputs })
1728    }
1729
1730    fn validate_inputs(&self, actual: &[ConcreteEinsumInputSpec], op: &'static str) -> Result<()> {
1731        if actual.len() != self.inputs.len() {
1732            return Err(Error::invalid_argument(
1733                op,
1734                "inputs",
1735                format!(
1736                    "prepared einsum expects {} inputs, got {}",
1737                    self.inputs.len(),
1738                    actual.len()
1739                ),
1740            ));
1741        }
1742
1743        for (expected, actual) in self.inputs.iter().zip(actual.iter()) {
1744            if expected.dtype != actual.dtype {
1745                return Err(Error::dtype_mismatch(op, expected.dtype, actual.dtype));
1746            }
1747            if expected.shape != actual.shape {
1748                return Err(Error::shape_mismatch(
1749                    op,
1750                    expected.shape.clone(),
1751                    actual.shape.clone(),
1752                ));
1753            }
1754        }
1755
1756        Ok(())
1757    }
1758}
1759
1760#[derive(Clone, Debug)]
1761struct ConcreteEinsumInputSpec {
1762    dtype: DType,
1763    shape: Vec<usize>,
1764}
1765
1766fn resolve_shapes(notation: &EinsumNotation, shapes: Vec<&[usize]>) -> Result<Subscripts> {
1767    resolve_einsum_notation(notation, &shapes)
1768}
1769
1770fn resolve_tensor_notation(inputs: &[&Tensor], notation: &EinsumNotation) -> Result<Subscripts> {
1771    resolve_shapes(
1772        notation,
1773        inputs.iter().map(|tensor| tensor.shape()).collect(),
1774    )
1775}
1776
1777fn resolve_typed_notation<T: TensorScalar>(
1778    inputs: &[&TypedTensor<T>],
1779    notation: &EinsumNotation,
1780) -> Result<Subscripts> {
1781    resolve_shapes(
1782        notation,
1783        inputs.iter().map(|tensor| tensor.shape()).collect(),
1784    )
1785}
1786
1787fn resolve_view_notation<'a, T: TensorScalar>(
1788    inputs: &[TypedTensorView<'a, T>],
1789    notation: &EinsumNotation,
1790) -> Result<Subscripts> {
1791    resolve_shapes(notation, inputs.iter().map(|view| view.shape()).collect())
1792}
1793
1794fn resolve_read_notation<'a>(
1795    inputs: &[TensorRead<'a>],
1796    notation: &EinsumNotation,
1797) -> Result<Subscripts> {
1798    resolve_shapes(notation, inputs.iter().map(|input| input.shape()).collect())
1799}
1800
1801fn input_specs(inputs: &[&Tensor]) -> Vec<ConcreteEinsumInputSpec> {
1802    inputs
1803        .iter()
1804        .map(|tensor| ConcreteEinsumInputSpec {
1805            dtype: tensor.dtype(),
1806            shape: tensor.shape().to_vec(),
1807        })
1808        .collect()
1809}
1810
1811fn typed_input_specs<T: TensorScalar>(inputs: &[&TypedTensor<T>]) -> Vec<ConcreteEinsumInputSpec> {
1812    inputs
1813        .iter()
1814        .map(|tensor| ConcreteEinsumInputSpec {
1815            dtype: T::dtype(),
1816            shape: tensor.shape().to_vec(),
1817        })
1818        .collect()
1819}
1820
1821fn read_input_specs(inputs: &[TensorRead<'_>]) -> Vec<ConcreteEinsumInputSpec> {
1822    inputs
1823        .iter()
1824        .map(|tensor| ConcreteEinsumInputSpec {
1825            dtype: tensor.dtype(),
1826            shape: tensor.shape().to_vec(),
1827        })
1828        .collect()
1829}
1830
1831fn typed_view_einsum_subscripts<T: TensorScalar>(
1832    session: &mut dyn BackendSession,
1833    inputs: &[TypedTensorView<'_, T>],
1834    subscripts: &Subscripts,
1835    op: &'static str,
1836) -> Result<TypedTensor<T>> {
1837    let reads: Vec<_> = inputs
1838        .iter()
1839        .cloned()
1840        .map(|view| TensorRead::from_view(T::tensor_view(view)))
1841        .collect();
1842    let plan =
1843        ConcreteEinsumPlan::prepare_subscripts_internal(read_input_specs(&reads), subscripts)?;
1844    let result = plan.execute_read(&reads, session)?;
1845    into_typed_result(result, op)
1846}
1847
1848fn tensor_einsum_into_subscripts(
1849    session: &mut dyn BackendSession,
1850    inputs: &[&Tensor],
1851    subscripts: &Subscripts,
1852    out: TensorWrite<'_>,
1853    op: &'static str,
1854) -> Result<()> {
1855    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(input_specs(inputs), subscripts)?;
1856    validate_output(&plan.inputs, &plan.tree, &out, op)?;
1857    plan.execute_into(inputs, session, out)
1858}
1859
1860fn typed_view_einsum_into_subscripts<T: TensorScalar>(
1861    session: &mut dyn BackendSession,
1862    inputs: &[TypedTensorView<'_, T>],
1863    subscripts: &Subscripts,
1864    out: TypedTensorWrite<'_, T>,
1865    op: &'static str,
1866) -> Result<()> {
1867    let reads: Vec<_> = inputs
1868        .iter()
1869        .cloned()
1870        .map(|view| TensorRead::from_view(T::tensor_view(view)))
1871        .collect();
1872    let plan =
1873        ConcreteEinsumPlan::prepare_subscripts_internal(read_input_specs(&reads), subscripts)?;
1874    let out = out.into_tensor_write();
1875    validate_output(&plan.inputs, &plan.tree, &out, op)?;
1876    plan.execute_read_into(&reads, session, out)
1877}
1878
1879fn typed_einsum_into_subscripts<T: TensorScalar>(
1880    session: &mut dyn BackendSession,
1881    inputs: &[&TypedTensor<T>],
1882    subscripts: &Subscripts,
1883    out: TypedTensorWrite<'_, T>,
1884    op: &'static str,
1885) -> Result<()> {
1886    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1887    let plan =
1888        ConcreteEinsumPlan::prepare_subscripts_internal(read_input_specs(&reads), subscripts)?;
1889    let out = out.into_tensor_write();
1890    validate_output(&plan.inputs, &plan.tree, &out, op)?;
1891    plan.execute_read_into(&reads, session, out)
1892}
1893
1894fn tensor_read_einsum_into_subscripts(
1895    session: &mut dyn BackendSession,
1896    inputs: &[TensorRead<'_>],
1897    subscripts: &Subscripts,
1898    out: TensorWrite<'_>,
1899    op: &'static str,
1900) -> Result<()> {
1901    let plan =
1902        ConcreteEinsumPlan::prepare_subscripts_internal(read_input_specs(inputs), subscripts)?;
1903    validate_output(&plan.inputs, &plan.tree, &out, op)?;
1904    plan.execute_read_into(inputs, session, out)
1905}
1906
1907fn validate_output(
1908    inputs: &[ConcreteEinsumInputSpec],
1909    tree: &ContractionTree,
1910    out: &TensorWrite<'_>,
1911    op: &'static str,
1912) -> Result<()> {
1913    let expected = output_spec(inputs, tree, op)?;
1914    if out.dtype() != expected.dtype {
1915        return Err(Error::dtype_mismatch(op, expected.dtype, out.dtype()));
1916    }
1917    if out.shape() != expected.shape.as_slice() {
1918        return Err(Error::shape_mismatch(
1919            op,
1920            out.shape().to_vec(),
1921            expected.shape.clone(),
1922        ));
1923    }
1924    Ok(())
1925}
1926
1927fn output_spec(
1928    inputs: &[ConcreteEinsumInputSpec],
1929    tree: &ContractionTree,
1930    op: &'static str,
1931) -> Result<ConcreteEinsumInputSpec> {
1932    let dtype = inputs
1933        .first()
1934        .ok_or_else(|| {
1935            Error::invalid_argument(op, "inputs", "einsum requires at least one input tensor")
1936        })?
1937        .dtype;
1938    for input in inputs {
1939        if input.dtype != dtype {
1940            return Err(Error::dtype_mismatch(op, dtype, input.dtype));
1941        }
1942    }
1943
1944    for (input, labels) in inputs.iter().zip(tree.subscripts.inputs.iter()) {
1945        if labels.len() != input.shape.len() {
1946            return Err(Error::rank_mismatch(op, labels.len(), input.shape.len()));
1947        }
1948    }
1949    let output_shape = tree.output_shape();
1950    if output_shape.len() != tree.subscripts.output.len() {
1951        return Err(Error::invalid_argument(
1952            op,
1953            "output labels",
1954            "an output label is missing from all inputs",
1955        ));
1956    }
1957    Ok(ConcreteEinsumInputSpec {
1958        dtype,
1959        shape: output_shape,
1960    })
1961}
1962
1963fn typed_einsum_subscripts<T: TensorScalar>(
1964    session: &mut dyn BackendSession,
1965    inputs: &[&TypedTensor<T>],
1966    subscripts: &Subscripts,
1967    op: &'static str,
1968) -> Result<TypedTensor<T>> {
1969    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1970    let plan =
1971        ConcreteEinsumPlan::prepare_subscripts_internal(read_input_specs(&reads), subscripts)?;
1972    let result = plan.execute_read(&reads, session)?;
1973    into_typed_result(result, op)
1974}
1975
1976pub(crate) fn into_typed_result<T: TensorScalar>(
1977    result: Tensor,
1978    op: &'static str,
1979) -> Result<TypedTensor<T>> {
1980    let actual = result.dtype();
1981    T::into_typed(result).map_err(|_| Error::dtype_mismatch(op, T::dtype(), actual))
1982}