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(¬ation, 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(¬ation, 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(¬ation, 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(¬ation, 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(¬ation, 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(¬ation, 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(¬ation, 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(¬ation, 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, ¬ation)
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, ¬ation)
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, ¬ation)
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}