pub struct Einsum2Plan<ID: AxisId> {
pub batch: Vec<ID>,
pub lo: Vec<ID>,
pub ro: Vec<ID>,
pub sum: Vec<ID>,
pub left_trace: Vec<ID>,
pub right_trace: Vec<ID>,
pub left_perm: Vec<usize>,
pub right_perm: Vec<usize>,
pub c_to_internal_perm: Vec<usize>,
}Expand description
Pre-computed execution plan for a binary einsum contraction.
Classifies axes into groups and precomputes the permutations needed to arrange operands for batched matrix multiplication.
Fields§
§batch: Vec<ID>Batch axes: present in A, B, and C.
lo: Vec<ID>Left-output axes: present in A and C, not in B.
ro: Vec<ID>Right-output axes: present in B and C, not in A.
sum: Vec<ID>Contraction axes: present in A and B, not in C.
left_trace: Vec<ID>Left trace axes: present only in A.
right_trace: Vec<ID>Right trace axes: present only in B.
left_perm: Vec<usize>Permutation to reorder A to [lo, sum, batch] after trace reduction.
right_perm: Vec<usize>Permutation to reorder B to [sum, ro, batch] after trace reduction.
c_to_internal_perm: Vec<usize>Permutation to reorder C from IC order to [lo, ro, batch] order.
Implementations§
Source§impl<ID: AxisId> Einsum2Plan<ID>
impl<ID: AxisId> Einsum2Plan<ID>
Sourcepub fn new(ia: &[ID], ib: &[ID], ic: &[ID]) -> Result<Self, EinsumError>
pub fn new(ia: &[ID], ib: &[ID], ic: &[ID]) -> Result<Self, EinsumError>
Build a plan from axis labels.
ia, ib, ic are the axis labels for A, B, C respectively.
Uses linear scans instead of hash collections for axis classification, which is faster for the small label sets typical in einsum contractions.
Sourcepub fn left_trace_indices(&self, ia: &[ID]) -> Vec<usize>
pub fn left_trace_indices(&self, ia: &[ID]) -> Vec<usize>
Get the indices of left_trace axes in the original ia array.
Sourcepub fn right_trace_indices(&self, ib: &[ID]) -> Vec<usize>
pub fn right_trace_indices(&self, ib: &[ID]) -> Vec<usize>
Get the indices of right_trace axes in the original ib array.
Sourcepub fn internal_to_c_perm(&self) -> Vec<usize>
pub fn internal_to_c_perm(&self) -> Vec<usize>
Get the inverse of c_to_internal_perm (maps \[batch, lo, ro\] back to IC order).
Trait Implementations§
Source§impl<ID: Clone + AxisId> Clone for Einsum2Plan<ID>
impl<ID: Clone + AxisId> Clone for Einsum2Plan<ID>
Source§fn clone(&self) -> Einsum2Plan<ID>
fn clone(&self) -> Einsum2Plan<ID>
1.0.0 · Source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
source. Read moreAuto Trait Implementations§
impl<ID> Freeze for Einsum2Plan<ID>
impl<ID> RefUnwindSafe for Einsum2Plan<ID>where
ID: RefUnwindSafe,
impl<ID> Send for Einsum2Plan<ID>where
ID: Send,
impl<ID> Sync for Einsum2Plan<ID>where
ID: Sync,
impl<ID> Unpin for Einsum2Plan<ID>where
ID: Unpin,
impl<ID> UnwindSafe for Einsum2Plan<ID>where
ID: UnwindSafe,
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
§impl<T> DistributionExt for Twhere
T: ?Sized,
impl<T> DistributionExt for Twhere
T: ?Sized,
fn rand<T>(&self, rng: &mut (impl Rng + ?Sized)) -> Twhere
Self: Distribution<T>,
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more