Enum EinSumOp 

pub enum EinSumOp<T> {
    TensorSource {
        tensor: CausalTensor<T>,
    },
    Contraction {
        lhs_axes: Vec<usize>,
        rhs_axes: Vec<usize>,
    },
    Reduction {
        axes: Vec<usize>,
    },
    MatMul,
    DotProd,
    Trace {
        axes1: usize,
        axes2: usize,
    },
    TensorProduct,
    ElementWiseProduct,
    Transpose {
        new_order: Vec<usize>,
    },
    DiagonalExtraction {
        axes1: usize,
        axes2: usize,
    },
    BatchMatMul,
}

Variants

TensorSource

Fields

tensor: CausalTensor<T>

Contraction

Fields

lhs_axes: Vec<usize>

rhs_axes: Vec<usize>

Reduction

Fields

axes: Vec<usize>

MatMul

Standard matrix multiplication (rank-2 tensors)

DotProd

Dot product (rank-1 tensors)

Trace

Trace over two specified axes

Fields

axes1: usize

axes2: usize

TensorProduct

Outer product

ElementWiseProduct

Hadamard product

Transpose

Fields

new_order: Vec<usize>

DiagonalExtraction

Fields

axes1: usize

axes2: usize

BatchMatMul

Batch matrix multiplication (rank-3 tensors)

Implementations

impl<T> EinSumOp<T>

pub fn tensor_source(tensor: CausalTensor<T>) -> EinSumAST<T>

Creates an EinSumAST leaf node representing a source CausalTensor.

This is the base case for the AST, where a tensor is directly provided as input.

Arguments

• tensor - The CausalTensor to be wrapped as a source node.

Returns

An EinSumAST<T> node containing the TensorSource operation.

pub fn contraction<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>( lhs: L, rhs: R, lhs_axes: Vec<usize>, rhs_axes: Vec<usize>, ) -> EinSumAST<T>

Creates an EinSumAST node for a generic tensor contraction operation.

This node represents a contraction between two CausalTensors along specified axes.

Arguments

• lhs - The left-hand side CausalTensor.
• rhs - The right-hand side CausalTensor.
• lhs_axes - A Vec<usize> specifying the axes of the lhs tensor to contract.
• rhs_axes - A Vec<usize> specifying the axes of the rhs tensor to contract.

Returns

An EinSumAST<T> node representing the contraction operation.

pub fn reduction<O: Into<CausalTensor<T>>>( operand: O, axes: Vec<usize>, ) -> EinSumAST<T>

Creates an EinSumAST node for a reduction operation.

This node represents summing elements of a CausalTensor along specified axes.

Arguments

• operand - The CausalTensor to perform the reduction on.
• axes - A Vec<usize> specifying the axes along which to sum.

Returns

An EinSumAST<T> node representing the reduction operation.

pub fn mat_mul<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>( lhs: L, rhs: R, ) -> EinSumAST<T>

Creates an EinSumAST node for a standard 2D matrix multiplication operation.

This node represents the multiplication of two rank-2 CausalTensors.

Arguments

• lhs - The left-hand side CausalTensor (matrix).
• rhs - The right-hand side CausalTensor (matrix).

Returns

An EinSumAST<T> node representing the matrix multiplication operation.

pub fn dot_prod<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>( lhs: L, rhs: R, ) -> EinSumAST<T>

Creates an EinSumAST node for a dot product operation.

This node represents the dot product of two rank-1 CausalTensors (vectors).

Arguments

• lhs - The left-hand side CausalTensor (vector).
• rhs - The right-hand side CausalTensor (vector).

Returns

An EinSumAST<T> node representing the dot product operation.

pub fn trace<O: Into<CausalTensor<T>>>( operand: O, axes1: usize, axes2: usize, ) -> EinSumAST<T>

Creates an EinSumAST node for a trace operation.

This node represents summing the diagonal elements of a CausalTensor over two specified axes.

Arguments

• operand - The CausalTensor to perform the trace on.
• axes1 - The first axis to trace over.
• axes2 - The second axis to trace over.

Returns

An EinSumAST<T> node representing the trace operation.

pub fn tensor_product<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>( lhs: L, rhs: R, ) -> EinSumAST<T>

Creates an EinSumAST node for an outer product (tensor product) operation.

This node represents the outer product of two CausalTensors.

Arguments

• lhs - The left-hand side CausalTensor.
• rhs - The right-hand side CausalTensor.

Returns

An EinSumAST<T> node representing the tensor product operation.

pub fn element_wise_product<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>( lhs: L, rhs: R, ) -> EinSumAST<T>

Creates an EinSumAST node for an element-wise product (Hadamard product) operation.

This node represents the element-wise multiplication of two CausalTensors.

Arguments

• lhs - The left-hand side CausalTensor.
• rhs - The right-hand side CausalTensor.

Returns

An EinSumAST<T> node representing the element-wise product operation.

pub fn transpose<O: Into<CausalTensor<T>>>( operand: O, new_order: Vec<usize>, ) -> EinSumAST<T>

Creates an EinSumAST node for a transpose operation.

This node represents permuting the axes of a CausalTensor according to a new order.

Arguments

• operand - The CausalTensor to transpose.
• new_order - A Vec<usize> specifying the new order of axes.

Returns

An EinSumAST<T> node representing the transpose operation.

pub fn diagonal_extraction<O: Into<CausalTensor<T>>>( operand: O, axes1: usize, axes2: usize, ) -> EinSumAST<T>

Creates an EinSumAST node for a diagonal extraction operation.

This node represents extracting the diagonal elements of a CausalTensor over two specified axes.

Arguments

• operand - The CausalTensor from which to extract the diagonal.
• axes1 - The first axis defining the 2D plane for diagonal extraction.
• axes2 - The second axis defining the 2D plane for diagonal extraction.

Returns

An EinSumAST<T> node representing the diagonal extraction operation.

pub fn batch_mat_mul<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>( lhs: L, rhs: R, ) -> EinSumAST<T>

Creates an EinSumAST node for a batch matrix multiplication operation.

This node represents performing matrix multiplication on batches of CausalTensors.

Arguments

• lhs - The left-hand side CausalTensor with a batch dimension.
• rhs - The right-hand side CausalTensor with a batch dimension.

Returns

An EinSumAST<T> node representing the batch matrix multiplication operation.