tenferro_einsum/
concrete.rs

1//! Public concrete tensor einsum extension API.
2
3use tenferro_tensor::{
4    DType, DotGeneralAccumulation, Error, Result, Tensor, TensorBackend, TensorRead, TensorScalar,
5    TensorWrite, TypedTensor, TypedTensorView,
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::{ContractionTree, EinsumSubscripts, Subscripts};
14
15const TENSOR_EINSUM_OP: &str = "TensorEinsumExt::einsum";
16const TENSOR_EINSUM_INTO_OP: &str = "TensorEinsumIntoExt::einsum_into";
17const TENSOR_READ_EINSUM_OP: &str = "TensorReadEinsumExt::einsum_read";
18const TENSOR_READ_EINSUM_INTO_OP: &str = "TensorReadEinsumIntoExt::einsum_read_into";
19const TYPED_TENSOR_EINSUM_OP: &str = "TypedTensorEinsumExt::einsum";
20const TYPED_TENSOR_EINSUM_INTO_OP: &str = "TypedTensorEinsumIntoExt::einsum_into";
21const PLAN_PREPARE_OP: &str = "ConcreteEinsumPlan::prepare";
22const PLAN_EXECUTE_OP: &str = "ConcreteEinsumPlan::execute";
23
24/// Backend-explicit einsum methods for dtype-erased concrete tensors.
25///
26/// Implementations are provided for slices and fixed-size arrays of
27/// [`Tensor`] references, so both `inputs.as_slice().einsum(...)` and
28/// `[&lhs, &rhs].einsum(...)` work.
29///
30/// # Examples
31///
32/// ```
33/// use tenferro_cpu::CpuBackend;
34/// use tenferro_einsum::TensorEinsumExt;
35/// use tenferro_tensor::Tensor;
36///
37/// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
38/// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
39/// let mut backend = CpuBackend::new();
40///
41/// let out = [&lhs, &rhs].einsum("ij,jk->ik", &mut backend)?;
42/// assert_eq!(out.shape(), &[2, 4]);
43/// # Ok::<(), tenferro_tensor::Error>(())
44/// ```
45pub trait TensorEinsumExt {
46    /// Execute an einsum from string notation.
47    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor>;
48
49    /// Execute an einsum from parsed integer-label subscripts.
50    fn einsum_subscripts<B: TensorBackend>(
51        &self,
52        subscripts: &EinsumSubscripts,
53        backend: &mut B,
54    ) -> Result<Tensor>;
55}
56
57impl TensorEinsumExt for [&Tensor] {
58    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
59        let subscripts = parse_subscripts(subscripts, TENSOR_EINSUM_OP)?;
60        eager_einsum_subscripts(backend, self, &subscripts)
61    }
62
63    fn einsum_subscripts<B: TensorBackend>(
64        &self,
65        subscripts: &EinsumSubscripts,
66        backend: &mut B,
67    ) -> Result<Tensor> {
68        let subscripts = Subscripts::from(subscripts);
69        eager_einsum_subscripts(backend, self, &subscripts)
70    }
71}
72
73impl<const N: usize> TensorEinsumExt for [&Tensor; N] {
74    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
75        self.as_slice().einsum(subscripts, backend)
76    }
77
78    fn einsum_subscripts<B: TensorBackend>(
79        &self,
80        subscripts: &EinsumSubscripts,
81        backend: &mut B,
82    ) -> Result<Tensor> {
83        self.as_slice().einsum_subscripts(subscripts, backend)
84    }
85}
86
87/// Backend-explicit preallocated-output einsum methods for dtype-erased tensors.
88pub trait TensorEinsumIntoExt {
89    /// Execute an einsum from string notation into caller-provided output.
90    fn einsum_into<B: TensorBackend>(
91        &self,
92        subscripts: &str,
93        backend: &mut B,
94        out: TensorWrite<'_>,
95    ) -> Result<()>;
96
97    /// Execute an einsum from parsed integer-label subscripts into caller-provided output.
98    fn einsum_into_subscripts<B: TensorBackend>(
99        &self,
100        subscripts: &EinsumSubscripts,
101        backend: &mut B,
102        out: TensorWrite<'_>,
103    ) -> Result<()>;
104}
105
106impl TensorEinsumIntoExt for [&Tensor] {
107    fn einsum_into<B: TensorBackend>(
108        &self,
109        subscripts: &str,
110        backend: &mut B,
111        out: TensorWrite<'_>,
112    ) -> Result<()> {
113        let subscripts = parse_subscripts(subscripts, TENSOR_EINSUM_INTO_OP)?;
114        tensor_einsum_into_subscripts(backend, self, &subscripts, out, TENSOR_EINSUM_INTO_OP)
115    }
116
117    fn einsum_into_subscripts<B: TensorBackend>(
118        &self,
119        subscripts: &EinsumSubscripts,
120        backend: &mut B,
121        out: TensorWrite<'_>,
122    ) -> Result<()> {
123        let subscripts = Subscripts::from(subscripts);
124        tensor_einsum_into_subscripts(backend, self, &subscripts, out, TENSOR_EINSUM_INTO_OP)
125    }
126}
127
128impl<const N: usize> TensorEinsumIntoExt for [&Tensor; N] {
129    fn einsum_into<B: TensorBackend>(
130        &self,
131        subscripts: &str,
132        backend: &mut B,
133        out: TensorWrite<'_>,
134    ) -> Result<()> {
135        self.as_slice().einsum_into(subscripts, backend, out)
136    }
137
138    fn einsum_into_subscripts<B: TensorBackend>(
139        &self,
140        subscripts: &EinsumSubscripts,
141        backend: &mut B,
142        out: TensorWrite<'_>,
143    ) -> Result<()> {
144        self.as_slice()
145            .einsum_into_subscripts(subscripts, backend, out)
146    }
147}
148
149/// Backend-explicit einsum methods for typed concrete tensors.
150///
151/// The result keeps the same scalar type as the inputs. Mixed dtypes should use
152/// [`TensorEinsumExt`] on dtype-erased [`Tensor`] values instead.
153///
154/// # Examples
155///
156/// ```
157/// use tenferro_cpu::CpuBackend;
158/// use tenferro_einsum::TypedTensorEinsumExt;
159/// use tenferro_tensor::TypedTensor;
160///
161/// let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
162/// let rhs = TypedTensor::<f64>::from_vec_col_major(vec![3, 4], vec![1.0; 12]).unwrap();
163/// let mut backend = CpuBackend::new();
164///
165/// let out = [&lhs, &rhs].einsum("ij,jk->ik", &mut backend)?;
166/// assert_eq!(out.shape(), &[2, 4]);
167/// # Ok::<(), tenferro_tensor::Error>(())
168/// ```
169pub trait TypedTensorEinsumExt<T: TensorScalar> {
170    /// Execute an einsum from string notation.
171    fn einsum<B: TensorBackend>(&self, subscripts: &str, backend: &mut B)
172        -> Result<TypedTensor<T>>;
173
174    /// Execute an einsum from parsed integer-label subscripts.
175    fn einsum_subscripts<B: TensorBackend>(
176        &self,
177        subscripts: &EinsumSubscripts,
178        backend: &mut B,
179    ) -> Result<TypedTensor<T>>;
180}
181
182impl<T: TensorScalar> TypedTensorEinsumExt<T> for [&TypedTensor<T>] {
183    fn einsum<B: TensorBackend>(
184        &self,
185        subscripts: &str,
186        backend: &mut B,
187    ) -> Result<TypedTensor<T>> {
188        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_EINSUM_OP)?;
189        typed_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
190    }
191
192    fn einsum_subscripts<B: TensorBackend>(
193        &self,
194        subscripts: &EinsumSubscripts,
195        backend: &mut B,
196    ) -> Result<TypedTensor<T>> {
197        let subscripts = Subscripts::from(subscripts);
198        typed_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
199    }
200}
201
202impl<T: TensorScalar, const N: usize> TypedTensorEinsumExt<T> for [&TypedTensor<T>; N] {
203    fn einsum<B: TensorBackend>(
204        &self,
205        subscripts: &str,
206        backend: &mut B,
207    ) -> Result<TypedTensor<T>> {
208        self.as_slice().einsum(subscripts, backend)
209    }
210
211    fn einsum_subscripts<B: TensorBackend>(
212        &self,
213        subscripts: &EinsumSubscripts,
214        backend: &mut B,
215    ) -> Result<TypedTensor<T>> {
216        self.as_slice().einsum_subscripts(subscripts, backend)
217    }
218}
219
220impl<'a, T: TensorScalar> TypedTensorEinsumExt<T> for [TypedTensorView<'a, T>] {
221    fn einsum<B: TensorBackend>(
222        &self,
223        subscripts: &str,
224        backend: &mut B,
225    ) -> Result<TypedTensor<T>> {
226        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_EINSUM_OP)?;
227        typed_view_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
228    }
229
230    fn einsum_subscripts<B: TensorBackend>(
231        &self,
232        subscripts: &EinsumSubscripts,
233        backend: &mut B,
234    ) -> Result<TypedTensor<T>> {
235        let subscripts = Subscripts::from(subscripts);
236        typed_view_einsum_subscripts(backend, self, &subscripts, TYPED_TENSOR_EINSUM_OP)
237    }
238}
239
240impl<'a, T: TensorScalar, const N: usize> TypedTensorEinsumExt<T> for [TypedTensorView<'a, T>; N] {
241    fn einsum<B: TensorBackend>(
242        &self,
243        subscripts: &str,
244        backend: &mut B,
245    ) -> Result<TypedTensor<T>> {
246        self.as_slice().einsum(subscripts, backend)
247    }
248
249    fn einsum_subscripts<B: TensorBackend>(
250        &self,
251        subscripts: &EinsumSubscripts,
252        backend: &mut B,
253    ) -> Result<TypedTensor<T>> {
254        self.as_slice().einsum_subscripts(subscripts, backend)
255    }
256}
257
258/// Backend-explicit preallocated-output einsum methods for typed concrete tensors.
259pub trait TypedTensorEinsumIntoExt<T: TensorScalar> {
260    /// Execute an einsum from string notation into caller-provided typed output.
261    fn einsum_into<B: TensorBackend>(
262        &self,
263        subscripts: &str,
264        backend: &mut B,
265        out: &mut TypedTensor<T>,
266    ) -> Result<()>;
267
268    /// Execute an einsum from parsed integer-label subscripts into caller-provided typed output.
269    fn einsum_into_subscripts<B: TensorBackend>(
270        &self,
271        subscripts: &EinsumSubscripts,
272        backend: &mut B,
273        out: &mut TypedTensor<T>,
274    ) -> Result<()>;
275}
276
277impl<T: TensorScalar> TypedTensorEinsumIntoExt<T> for [&TypedTensor<T>] {
278    fn einsum_into<B: TensorBackend>(
279        &self,
280        subscripts: &str,
281        backend: &mut B,
282        out: &mut TypedTensor<T>,
283    ) -> Result<()> {
284        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_EINSUM_INTO_OP)?;
285        typed_einsum_into_subscripts(backend, self, &subscripts, out, TYPED_TENSOR_EINSUM_INTO_OP)
286    }
287
288    fn einsum_into_subscripts<B: TensorBackend>(
289        &self,
290        subscripts: &EinsumSubscripts,
291        backend: &mut B,
292        out: &mut TypedTensor<T>,
293    ) -> Result<()> {
294        let subscripts = Subscripts::from(subscripts);
295        typed_einsum_into_subscripts(backend, self, &subscripts, out, TYPED_TENSOR_EINSUM_INTO_OP)
296    }
297}
298
299impl<T: TensorScalar, const N: usize> TypedTensorEinsumIntoExt<T> for [&TypedTensor<T>; N] {
300    fn einsum_into<B: TensorBackend>(
301        &self,
302        subscripts: &str,
303        backend: &mut B,
304        out: &mut TypedTensor<T>,
305    ) -> Result<()> {
306        self.as_slice().einsum_into(subscripts, backend, out)
307    }
308
309    fn einsum_into_subscripts<B: TensorBackend>(
310        &self,
311        subscripts: &EinsumSubscripts,
312        backend: &mut B,
313        out: &mut TypedTensor<T>,
314    ) -> Result<()> {
315        self.as_slice()
316            .einsum_into_subscripts(subscripts, backend, out)
317    }
318}
319
320impl<'a, T: TensorScalar> TypedTensorEinsumIntoExt<T> for [TypedTensorView<'a, T>] {
321    fn einsum_into<B: TensorBackend>(
322        &self,
323        subscripts: &str,
324        backend: &mut B,
325        out: &mut TypedTensor<T>,
326    ) -> Result<()> {
327        let subscripts = parse_subscripts(subscripts, TYPED_TENSOR_EINSUM_INTO_OP)?;
328        typed_view_einsum_into_subscripts(
329            backend,
330            self,
331            &subscripts,
332            out,
333            TYPED_TENSOR_EINSUM_INTO_OP,
334        )
335    }
336
337    fn einsum_into_subscripts<B: TensorBackend>(
338        &self,
339        subscripts: &EinsumSubscripts,
340        backend: &mut B,
341        out: &mut TypedTensor<T>,
342    ) -> Result<()> {
343        let subscripts = Subscripts::from(subscripts);
344        typed_view_einsum_into_subscripts(
345            backend,
346            self,
347            &subscripts,
348            out,
349            TYPED_TENSOR_EINSUM_INTO_OP,
350        )
351    }
352}
353
354impl<'a, T: TensorScalar, const N: usize> TypedTensorEinsumIntoExt<T>
355    for [TypedTensorView<'a, T>; N]
356{
357    fn einsum_into<B: TensorBackend>(
358        &self,
359        subscripts: &str,
360        backend: &mut B,
361        out: &mut TypedTensor<T>,
362    ) -> Result<()> {
363        self.as_slice().einsum_into(subscripts, backend, out)
364    }
365
366    fn einsum_into_subscripts<B: TensorBackend>(
367        &self,
368        subscripts: &EinsumSubscripts,
369        backend: &mut B,
370        out: &mut TypedTensor<T>,
371    ) -> Result<()> {
372        self.as_slice()
373            .einsum_into_subscripts(subscripts, backend, out)
374    }
375}
376
377/// Backend-explicit einsum methods for [`TensorRead`] inputs.
378///
379/// Use this surface when an input is a borrowed tensor view rather than an
380/// owned compact [`Tensor`]. The `_read` suffix follows the repository-wide
381/// convention for APIs that explicitly accept read-oriented borrowed inputs.
382///
383/// # Examples
384///
385/// ```
386/// use tenferro_cpu::CpuBackend;
387/// use tenferro_einsum::TensorReadEinsumExt;
388/// use tenferro_tensor::{Tensor, TensorRead, TensorView};
389///
390/// let shape = [2, 3];
391/// let data = [1.0_f64; 6];
392/// let rhs = Tensor::from_vec_col_major(vec![3], vec![1.0_f64; 3]).unwrap();
393/// let inputs = [
394///     TensorRead::from_view(TensorView::f64(&shape, &data)?),
395///     TensorRead::from_tensor(&rhs),
396/// ];
397/// let mut backend = CpuBackend::new();
398///
399/// let out = inputs.einsum_read("ij,j->i", &mut backend)?;
400/// assert_eq!(out.shape(), &[2]);
401/// # Ok::<(), tenferro_tensor::Error>(())
402/// ```
403pub trait TensorReadEinsumExt {
404    /// Execute an einsum from string notation over read-only tensor inputs.
405    fn einsum_read<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor>;
406
407    /// Execute an einsum from parsed integer-label subscripts over read-only
408    /// tensor inputs.
409    fn einsum_read_subscripts<B: TensorBackend>(
410        &self,
411        subscripts: &EinsumSubscripts,
412        backend: &mut B,
413    ) -> Result<Tensor>;
414}
415
416impl<'a> TensorReadEinsumExt for [TensorRead<'a>] {
417    fn einsum_read<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
418        let subscripts = parse_subscripts(subscripts, TENSOR_READ_EINSUM_OP)?;
419        eager_einsum_read_subscripts(backend, self, &subscripts)
420    }
421
422    fn einsum_read_subscripts<B: TensorBackend>(
423        &self,
424        subscripts: &EinsumSubscripts,
425        backend: &mut B,
426    ) -> Result<Tensor> {
427        let subscripts = Subscripts::from(subscripts);
428        eager_einsum_read_subscripts(backend, self, &subscripts)
429    }
430}
431
432impl<'a, const N: usize> TensorReadEinsumExt for [TensorRead<'a>; N] {
433    fn einsum_read<B: TensorBackend>(&self, subscripts: &str, backend: &mut B) -> Result<Tensor> {
434        self.as_slice().einsum_read(subscripts, backend)
435    }
436
437    fn einsum_read_subscripts<B: TensorBackend>(
438        &self,
439        subscripts: &EinsumSubscripts,
440        backend: &mut B,
441    ) -> Result<Tensor> {
442        self.as_slice().einsum_read_subscripts(subscripts, backend)
443    }
444}
445
446/// Backend-explicit preallocated-output einsum methods for [`TensorRead`] inputs.
447pub trait TensorReadEinsumIntoExt {
448    /// Execute an einsum from string notation over read-only inputs into caller-provided output.
449    fn einsum_read_into<B: TensorBackend>(
450        &self,
451        subscripts: &str,
452        backend: &mut B,
453        out: TensorWrite<'_>,
454    ) -> Result<()>;
455
456    /// Execute an einsum from parsed integer-label subscripts over read-only inputs into output.
457    fn einsum_read_into_subscripts<B: TensorBackend>(
458        &self,
459        subscripts: &EinsumSubscripts,
460        backend: &mut B,
461        out: TensorWrite<'_>,
462    ) -> Result<()>;
463}
464
465impl<'a> TensorReadEinsumIntoExt for [TensorRead<'a>] {
466    fn einsum_read_into<B: TensorBackend>(
467        &self,
468        subscripts: &str,
469        backend: &mut B,
470        out: TensorWrite<'_>,
471    ) -> Result<()> {
472        let subscripts = parse_subscripts(subscripts, TENSOR_READ_EINSUM_INTO_OP)?;
473        tensor_read_einsum_into_subscripts(
474            backend,
475            self,
476            &subscripts,
477            out,
478            TENSOR_READ_EINSUM_INTO_OP,
479        )
480    }
481
482    fn einsum_read_into_subscripts<B: TensorBackend>(
483        &self,
484        subscripts: &EinsumSubscripts,
485        backend: &mut B,
486        out: TensorWrite<'_>,
487    ) -> Result<()> {
488        let subscripts = Subscripts::from(subscripts);
489        tensor_read_einsum_into_subscripts(
490            backend,
491            self,
492            &subscripts,
493            out,
494            TENSOR_READ_EINSUM_INTO_OP,
495        )
496    }
497}
498
499impl<'a, const N: usize> TensorReadEinsumIntoExt for [TensorRead<'a>; N] {
500    fn einsum_read_into<B: TensorBackend>(
501        &self,
502        subscripts: &str,
503        backend: &mut B,
504        out: TensorWrite<'_>,
505    ) -> Result<()> {
506        self.as_slice().einsum_read_into(subscripts, backend, out)
507    }
508
509    fn einsum_read_into_subscripts<B: TensorBackend>(
510        &self,
511        subscripts: &EinsumSubscripts,
512        backend: &mut B,
513        out: TensorWrite<'_>,
514    ) -> Result<()> {
515        self.as_slice()
516            .einsum_read_into_subscripts(subscripts, backend, out)
517    }
518}
519
520/// Prepared concrete einsum plan for repeated executions with fixed input
521/// dtype and shape metadata.
522///
523/// Preparing a plan parses and optimizes the contraction tree once. Execution
524/// validates the later inputs against the prepared dtype and shape contract,
525/// then runs the stored tree without re-planning.
526///
527/// # Examples
528///
529/// ```
530/// use tenferro_cpu::CpuBackend;
531/// use tenferro_einsum::ConcreteEinsumPlan;
532/// use tenferro_tensor::Tensor;
533///
534/// let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
535/// let rhs = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
536/// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik")?;
537///
538/// let mut backend = CpuBackend::new();
539/// let out = plan.execute([&lhs, &rhs], &mut backend)?;
540/// assert_eq!(out.shape(), &[2, 4]);
541/// # Ok::<(), tenferro_tensor::Error>(())
542/// ```
543#[derive(Debug)]
544pub struct ConcreteEinsumPlan {
545    tree: ContractionTree,
546    inputs: Vec<ConcreteEinsumInputSpec>,
547}
548
549impl ConcreteEinsumPlan {
550    /// Prepare a plan from dtype-erased concrete tensor inputs and string
551    /// notation.
552    pub fn prepare<'a, I>(inputs: I, subscripts: &str) -> Result<Self>
553    where
554        I: AsRef<[&'a Tensor]>,
555    {
556        let subscripts = parse_subscripts(subscripts, PLAN_PREPARE_OP)?;
557        Self::prepare_subscripts_internal(input_specs(inputs.as_ref()), &subscripts)
558    }
559
560    /// Prepare a plan from dtype-erased concrete tensor inputs and parsed
561    /// integer-label subscripts.
562    pub fn prepare_subscripts<'a, I>(inputs: I, subscripts: &EinsumSubscripts) -> Result<Self>
563    where
564        I: AsRef<[&'a Tensor]>,
565    {
566        let subscripts = Subscripts::from(subscripts);
567        Self::prepare_subscripts_internal(input_specs(inputs.as_ref()), &subscripts)
568    }
569
570    /// Prepare a plan from typed concrete tensor inputs and string notation.
571    pub fn prepare_typed<'a, T, I>(inputs: I, subscripts: &str) -> Result<Self>
572    where
573        T: TensorScalar,
574        I: AsRef<[&'a TypedTensor<T>]>,
575    {
576        let subscripts = parse_subscripts(subscripts, PLAN_PREPARE_OP)?;
577        Self::prepare_subscripts_internal(typed_input_specs(inputs.as_ref()), &subscripts)
578    }
579
580    /// Prepare a plan from typed concrete tensor inputs and parsed integer-label
581    /// subscripts.
582    pub fn prepare_typed_subscripts<'a, T, I>(
583        inputs: I,
584        subscripts: &EinsumSubscripts,
585    ) -> Result<Self>
586    where
587        T: TensorScalar,
588        I: AsRef<[&'a TypedTensor<T>]>,
589    {
590        let subscripts = Subscripts::from(subscripts);
591        Self::prepare_subscripts_internal(typed_input_specs(inputs.as_ref()), &subscripts)
592    }
593
594    /// Prepare a plan from read-only tensor inputs and string notation.
595    pub fn prepare_read<'a, I>(inputs: I, subscripts: &str) -> Result<Self>
596    where
597        I: AsRef<[TensorRead<'a>]>,
598    {
599        let subscripts = parse_subscripts(subscripts, PLAN_PREPARE_OP)?;
600        Self::prepare_subscripts_internal(read_input_specs(inputs.as_ref()), &subscripts)
601    }
602
603    /// Prepare a plan from read-only tensor inputs and parsed integer-label
604    /// subscripts.
605    pub fn prepare_read_subscripts<'a, I>(inputs: I, subscripts: &EinsumSubscripts) -> Result<Self>
606    where
607        I: AsRef<[TensorRead<'a>]>,
608    {
609        let subscripts = Subscripts::from(subscripts);
610        Self::prepare_subscripts_internal(read_input_specs(inputs.as_ref()), &subscripts)
611    }
612
613    /// Execute this plan on dtype-erased concrete tensor inputs.
614    pub fn execute<'a, I, B>(&self, inputs: I, backend: &mut B) -> Result<Tensor>
615    where
616        I: AsRef<[&'a Tensor]>,
617        B: TensorBackend,
618    {
619        let inputs = inputs.as_ref();
620        self.validate_inputs(&input_specs(inputs), PLAN_EXECUTE_OP)?;
621        backend.with_backend_session(|exec| eager_einsum_exec(exec, inputs, &self.tree))
622    }
623
624    /// Execute this plan on typed concrete tensor inputs.
625    pub fn execute_typed<'a, T, I, B>(&self, inputs: I, backend: &mut B) -> Result<TypedTensor<T>>
626    where
627        T: TensorScalar,
628        I: AsRef<[&'a TypedTensor<T>]>,
629        B: TensorBackend,
630    {
631        let inputs = inputs.as_ref();
632        self.validate_inputs(&typed_input_specs(inputs), PLAN_EXECUTE_OP)?;
633        let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
634        let result = backend
635            .with_backend_session(|exec| eager_einsum_exec_read(exec, &reads, &self.tree))?;
636        into_typed_result(result, PLAN_EXECUTE_OP)
637    }
638
639    /// Execute this plan on read-only tensor inputs.
640    pub fn execute_read<'a, I, B>(&self, inputs: I, backend: &mut B) -> Result<Tensor>
641    where
642        I: AsRef<[TensorRead<'a>]>,
643        B: TensorBackend,
644    {
645        let inputs = inputs.as_ref();
646        self.validate_inputs(&read_input_specs(inputs), PLAN_EXECUTE_OP)?;
647        backend.with_backend_session(|exec| eager_einsum_exec_read(exec, inputs, &self.tree))
648    }
649
650    /// Execute this plan on dtype-erased concrete tensor inputs into caller-provided output.
651    pub fn execute_into<'a, I, B>(
652        &self,
653        inputs: I,
654        backend: &mut B,
655        out: TensorWrite<'_>,
656    ) -> Result<()>
657    where
658        I: AsRef<[&'a Tensor]>,
659        B: TensorBackend,
660    {
661        let inputs = inputs.as_ref();
662        let specs = input_specs(inputs);
663        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
664        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
665        let reads: Vec<_> = inputs
666            .iter()
667            .map(|tensor| TensorRead::from_tensor(tensor))
668            .collect();
669        backend
670            .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &self.tree, out))
671    }
672
673    /// Execute this plan on typed concrete tensor inputs into caller-provided output.
674    pub fn execute_typed_into<'a, T, I, B>(
675        &self,
676        inputs: I,
677        backend: &mut B,
678        out: &mut TypedTensor<T>,
679    ) -> Result<()>
680    where
681        T: TensorScalar,
682        I: AsRef<[&'a TypedTensor<T>]>,
683        B: TensorBackend,
684    {
685        let inputs = inputs.as_ref();
686        let specs = typed_input_specs(inputs);
687        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
688        let out = T::tensor_write(out);
689        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
690        let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
691        backend
692            .with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &self.tree, out))
693    }
694
695    /// Execute this plan on read-only tensor inputs into caller-provided output.
696    pub fn execute_read_into<'a, I, B>(
697        &self,
698        inputs: I,
699        backend: &mut B,
700        out: TensorWrite<'_>,
701    ) -> Result<()>
702    where
703        I: AsRef<[TensorRead<'a>]>,
704        B: TensorBackend,
705    {
706        let inputs = inputs.as_ref();
707        let specs = read_input_specs(inputs);
708        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
709        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
710        backend
711            .with_backend_session(|exec| eager_einsum_exec_read_into(exec, inputs, &self.tree, out))
712    }
713
714    /// Execute this plan on read-only inputs with scaled output accumulation.
715    ///
716    /// `accumulation` follows the dot-general contract:
717    /// `out = alpha * einsum(inputs) + beta * out`.
718    ///
719    /// # Examples
720    ///
721    /// ```rust
722    /// use tenferro_cpu::CpuBackend;
723    /// use tenferro_einsum::ConcreteEinsumPlan;
724    /// use tenferro_tensor::{
725    ///     DotGeneralAccumulation, DType, Tensor, TensorRead, TensorWrite,
726    /// };
727    ///
728    /// let lhs = Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?;
729    /// let rhs = Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?;
730    /// let mut out = Tensor::from_vec_col_major(vec![], vec![1.0_f64])?;
731    /// let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "i,i->")?;
732    /// let mut backend = CpuBackend::new();
733    /// plan.execute_read_into_accum(
734    ///     [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
735    ///     &mut backend,
736    ///     DotGeneralAccumulation::add_to(DType::F64)?,
737    ///     TensorWrite::from_tensor(&mut out),
738    /// )?;
739    /// assert_eq!(out.as_slice::<f64>()?, &[7.0]);
740    /// # Ok::<(), tenferro_tensor::Error>(())
741    /// ```
742    pub fn execute_read_into_accum<'a, I, B>(
743        &self,
744        inputs: I,
745        backend: &mut B,
746        accumulation: DotGeneralAccumulation,
747        out: TensorWrite<'_>,
748    ) -> Result<()>
749    where
750        I: AsRef<[TensorRead<'a>]>,
751        B: TensorBackend,
752    {
753        let inputs = inputs.as_ref();
754        let specs = read_input_specs(inputs);
755        self.validate_inputs(&specs, PLAN_EXECUTE_OP)?;
756        validate_output(&self.inputs, &self.tree, &out, PLAN_EXECUTE_OP)?;
757        backend.with_backend_session(|exec| {
758            eager_einsum_exec_read_into_accum(exec, inputs, &self.tree, accumulation, out)
759        })
760    }
761
762    fn prepare_subscripts_internal(
763        inputs: Vec<ConcreteEinsumInputSpec>,
764        subscripts: &Subscripts,
765    ) -> Result<Self> {
766        let shapes: Vec<&[usize]> = inputs.iter().map(|input| input.shape.as_slice()).collect();
767        let tree = plan_subscripts(subscripts, &shapes)?;
768        Ok(Self { tree, inputs })
769    }
770
771    fn validate_inputs(&self, actual: &[ConcreteEinsumInputSpec], op: &'static str) -> Result<()> {
772        if actual.len() != self.inputs.len() {
773            return Err(Error::InvalidConfig {
774                op,
775                message: format!(
776                    "prepared einsum expects {} inputs, got {}",
777                    self.inputs.len(),
778                    actual.len()
779                ),
780            });
781        }
782
783        for (expected, actual) in self.inputs.iter().zip(actual.iter()) {
784            if expected.dtype != actual.dtype {
785                return Err(Error::DTypeMismatch {
786                    op,
787                    lhs: expected.dtype,
788                    rhs: actual.dtype,
789                });
790            }
791            if expected.shape != actual.shape {
792                return Err(Error::ShapeMismatch {
793                    op,
794                    lhs: expected.shape.clone(),
795                    rhs: actual.shape.clone(),
796                });
797            }
798        }
799
800        Ok(())
801    }
802}
803
804#[derive(Clone, Debug)]
805struct ConcreteEinsumInputSpec {
806    dtype: DType,
807    shape: Vec<usize>,
808}
809
810fn parse_subscripts(subscripts: &str, op: &'static str) -> Result<Subscripts> {
811    Subscripts::parse(subscripts).map_err(|err| Error::InvalidConfig {
812        op,
813        message: format!("invalid subscripts: {err}"),
814    })
815}
816
817fn input_specs(inputs: &[&Tensor]) -> Vec<ConcreteEinsumInputSpec> {
818    inputs
819        .iter()
820        .map(|tensor| ConcreteEinsumInputSpec {
821            dtype: tensor.dtype(),
822            shape: tensor.shape().to_vec(),
823        })
824        .collect()
825}
826
827fn typed_input_specs<T: TensorScalar>(inputs: &[&TypedTensor<T>]) -> Vec<ConcreteEinsumInputSpec> {
828    inputs
829        .iter()
830        .map(|tensor| ConcreteEinsumInputSpec {
831            dtype: T::dtype(),
832            shape: tensor.shape().to_vec(),
833        })
834        .collect()
835}
836
837fn typed_view_input_specs<T: TensorScalar>(
838    inputs: &[TypedTensorView<'_, T>],
839) -> Vec<ConcreteEinsumInputSpec> {
840    inputs
841        .iter()
842        .map(|tensor| ConcreteEinsumInputSpec {
843            dtype: T::dtype(),
844            shape: tensor.shape().to_vec(),
845        })
846        .collect()
847}
848
849fn read_input_specs(inputs: &[TensorRead<'_>]) -> Vec<ConcreteEinsumInputSpec> {
850    inputs
851        .iter()
852        .map(|tensor| ConcreteEinsumInputSpec {
853            dtype: tensor.dtype(),
854            shape: tensor.shape().to_vec(),
855        })
856        .collect()
857}
858
859fn typed_view_einsum_subscripts<T: TensorScalar>(
860    backend: &mut impl TensorBackend,
861    inputs: &[TypedTensorView<'_, T>],
862    subscripts: &Subscripts,
863    op: &'static str,
864) -> Result<TypedTensor<T>> {
865    let reads: Vec<_> = inputs
866        .iter()
867        .cloned()
868        .map(|view| TensorRead::from_view(T::tensor_view(view)))
869        .collect();
870    let result = eager_einsum_read_subscripts(backend, &reads, subscripts)?;
871    into_typed_result(result, op)
872}
873
874fn tensor_einsum_into_subscripts(
875    backend: &mut impl TensorBackend,
876    inputs: &[&Tensor],
877    subscripts: &Subscripts,
878    out: TensorWrite<'_>,
879    op: &'static str,
880) -> Result<()> {
881    let specs = input_specs(inputs);
882    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
883    validate_output(&specs, &plan.tree, &out, op)?;
884    let reads: Vec<_> = inputs
885        .iter()
886        .map(|tensor| TensorRead::from_tensor(tensor))
887        .collect();
888    backend.with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &plan.tree, out))
889}
890
891fn typed_view_einsum_into_subscripts<T: TensorScalar>(
892    backend: &mut impl TensorBackend,
893    inputs: &[TypedTensorView<'_, T>],
894    subscripts: &Subscripts,
895    out: &mut TypedTensor<T>,
896    op: &'static str,
897) -> Result<()> {
898    let specs = typed_view_input_specs(inputs);
899    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
900    let out = T::tensor_write(out);
901    validate_output(&specs, &plan.tree, &out, op)?;
902    let reads: Vec<_> = inputs
903        .iter()
904        .cloned()
905        .map(|view| TensorRead::from_view(T::tensor_view(view)))
906        .collect();
907    backend.with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &plan.tree, out))
908}
909
910fn typed_einsum_into_subscripts<T: TensorScalar>(
911    backend: &mut impl TensorBackend,
912    inputs: &[&TypedTensor<T>],
913    subscripts: &Subscripts,
914    out: &mut TypedTensor<T>,
915    op: &'static str,
916) -> Result<()> {
917    let specs = typed_input_specs(inputs);
918    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
919    let out = T::tensor_write(out);
920    validate_output(&specs, &plan.tree, &out, op)?;
921    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
922    backend.with_backend_session(|exec| eager_einsum_exec_read_into(exec, &reads, &plan.tree, out))
923}
924
925fn tensor_read_einsum_into_subscripts(
926    backend: &mut impl TensorBackend,
927    inputs: &[TensorRead<'_>],
928    subscripts: &Subscripts,
929    out: TensorWrite<'_>,
930    op: &'static str,
931) -> Result<()> {
932    let specs = read_input_specs(inputs);
933    let plan = ConcreteEinsumPlan::prepare_subscripts_internal(specs.clone(), subscripts)?;
934    validate_output(&specs, &plan.tree, &out, op)?;
935    backend.with_backend_session(|exec| eager_einsum_exec_read_into(exec, inputs, &plan.tree, out))
936}
937
938fn validate_output(
939    inputs: &[ConcreteEinsumInputSpec],
940    tree: &ContractionTree,
941    out: &TensorWrite<'_>,
942    op: &'static str,
943) -> Result<()> {
944    let expected = output_spec(inputs, tree, op)?;
945    if out.dtype() != expected.dtype {
946        return Err(Error::DTypeMismatch {
947            op,
948            lhs: out.dtype(),
949            rhs: expected.dtype,
950        });
951    }
952    if out.shape() != expected.shape.as_slice() {
953        return Err(Error::ShapeMismatch {
954            op,
955            lhs: out.shape().to_vec(),
956            rhs: expected.shape,
957        });
958    }
959    Ok(())
960}
961
962fn output_spec(
963    inputs: &[ConcreteEinsumInputSpec],
964    tree: &ContractionTree,
965    op: &'static str,
966) -> Result<ConcreteEinsumInputSpec> {
967    let dtype = inputs
968        .first()
969        .ok_or_else(|| Error::InvalidConfig {
970            op,
971            message: "einsum requires at least one input tensor".to_string(),
972        })?
973        .dtype;
974    for input in inputs {
975        if input.dtype != dtype {
976            return Err(Error::DTypeMismatch {
977                op,
978                lhs: dtype,
979                rhs: input.dtype,
980            });
981        }
982    }
983
984    let mut output_shape = Vec::with_capacity(tree.subscripts.output.len());
985    for &label in &tree.subscripts.output {
986        let mut found = None;
987        for (input, labels) in inputs.iter().zip(tree.subscripts.inputs.iter()) {
988            if labels.len() != input.shape.len() {
989                return Err(Error::RankMismatch {
990                    op,
991                    expected: labels.len(),
992                    actual: input.shape.len(),
993                });
994            }
995            if let Some(axis) = labels.iter().position(|candidate| *candidate == label) {
996                found = Some(input.shape[axis]);
997                break;
998            }
999        }
1000        let Some(extent) = found else {
1001            return Err(Error::InvalidConfig {
1002                op,
1003                message: format!("output label {label} is missing from inputs"),
1004            });
1005        };
1006        output_shape.push(extent);
1007    }
1008    Ok(ConcreteEinsumInputSpec {
1009        dtype,
1010        shape: output_shape,
1011    })
1012}
1013
1014fn typed_einsum_subscripts<T: TensorScalar>(
1015    backend: &mut impl TensorBackend,
1016    inputs: &[&TypedTensor<T>],
1017    subscripts: &Subscripts,
1018    op: &'static str,
1019) -> Result<TypedTensor<T>> {
1020    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
1021    let result = eager_einsum_read_subscripts(backend, &reads, subscripts)?;
1022    into_typed_result(result, op)
1023}
1024
1025pub(crate) fn into_typed_result<T: TensorScalar>(
1026    result: Tensor,
1027    op: &'static str,
1028) -> Result<TypedTensor<T>> {
1029    let actual = result.dtype();
1030    T::into_typed(result).map_err(|_| Error::DTypeMismatch {
1031        op,
1032        lhs: T::dtype(),
1033        rhs: actual,
1034    })
1035}
tenferro_einsum/concrete.rs

tenferro_einsum/
concrete.rs
