boundlab.linearop.EinsumOp

class boundlab.linearop.EinsumOp

Bases: LinearOp

A linear operator defined by an Einstein summation with a fixed tensor.

Depending on input_dims and output_dims, this can express contraction (dot-product-like behavior), Hadamard-style elementwise multiplication, and dimension expansion.

Methods

__init__	Initialize an EinsumOp.
abs	Return a LinearOp representing the element-wise absolute value of this EinsumOp.
add_conditions	Return a new EinsumOp with additional (input_dim, output_dim) pairs multiplied together.
backward	Apply the transposed linear function to an input tensor.
einsum_op	Materialize this LinearOp as an explicit Jacobian tensor.
force_jacobian	Materialize Jacobian via batched forward/backward application.
forward	Apply the original linear function to an input tensor.
from_full	Create an EinsumOp that fully contracts over the specified input dimension.
from_hardmard	Create an EinsumOp for Hadamard-style multiplication with tensor.
from_scalar	Convert a ScalarOp to an EinsumOp.
is_full	Check if this EinsumOp fully contracts all input dimensions (output_dims & input_dims == ø).
is_hardmard	Check whether this EinsumOp performs no contraction over input dims.
is_non_expanding	Check if this EinsumOp doesn't introduce new dimensions (output_dims - input_dims == ø).
is_tensordot	Check if this EinsumOp is effectively a tensordot (no elementwise multiplication).
jacobian	Return an explicit Jacobian for full-layout EinsumOp instances.
jacobian_scatter	Add this operator's Jacobian contribution into an existing tensor.
norm_input	Return a LinearOp that computes the norm over the input dimensions of this EinsumOp.
norm_output	Return a LinearOp that computes the norm over the output dimensions of this EinsumOp.
permute_for_input
permute_for_output
purify_with
remove_conditions	Return a new EinsumOp with some (input_dim, output_dim) pairs no longer multiplied together.
squeeze_input	Drop input_dims[idx] whose corresponding tensor dim has size 1.
squeeze_output	Drop output_dims[idx] whose corresponding tensor dim has size 1.
sum_input	Return a LinearOp that sums over the input dimensions of this EinsumOp.
sum_output	Return a LinearOp that sums over the output dimensions of this EinsumOp.
unsqueeze_input	Insert a new size-1 input dim at position idx.
unsqueeze_output	Insert a new size-1 output dim at position idx.
vbackward	Apply the transposed linear function to an input tensor, supporting additional leading dimensions for batching.
vforward	Apply the original linear function to an input tensor, supporting additional trailing dimensions for batching.
with_tensor	Return a new EinsumOp with the same input/output dims but a different tensor.

__init__(tensor, input_dims, output_dims, name=None)

Initialize an EinsumOp.

This operator behaves like contraction on input_dims - output_dims, elementwise multiplication on input_dims & output_dims, and expansion on output_dims - input_dims.

Parameters:

• tensor (torch.Tensor) – The fixed tensor used in the Einstein summation.

• input_dims (list[int]) – A list of dimensions of tensor that correspond to the input tensor dimensions.

• output_dims (list[int]) – A list of dimensions of tensor that correspond to the output tensor dimensions.

• name – Optional name for display purposes.

forward(x)

Apply the original linear function to an input tensor.

backward(grad)

Apply the transposed linear function to an input tensor.

vforward(x)

Apply the original linear function to an input tensor, supporting additional trailing dimensions for batching.

vbackward(grad)

Apply the transposed linear function to an input tensor, supporting additional leading dimensions for batching.

is_full()

Check if this EinsumOp fully contracts all input dimensions (output_dims & input_dims == ø).

is_hardmard()

Check whether this EinsumOp performs no contraction over input dims.

Note

The method name keeps the historical hardmard spelling for backward compatibility.

is_non_expanding()

Check if this EinsumOp doesn’t introduce new dimensions (output_dims - input_dims == ø).

is_tensordot()

Check if this EinsumOp is effectively a tensordot (no elementwise multiplication).

static from_hardmard(tensor, n_input_dims=None, name=None)

Create an EinsumOp for Hadamard-style multiplication with tensor.

Note

The constructor name keeps the historical hardmard spelling for backward compatibility.

static from_full(tensor, input_dim, name=None)

Create an EinsumOp that fully contracts over the specified input dimension.

__mul__(scalar)

Scale the EinsumOp by a scalar.

squeeze_input(idx)

Drop input_dims[idx] whose corresponding tensor dim has size 1.

squeeze_output(idx)

Drop output_dims[idx] whose corresponding tensor dim has size 1.

unsqueeze_input(idx)

Insert a new size-1 input dim at position idx.

unsqueeze_output(idx)

Insert a new size-1 output dim at position idx.

__add__(other)

Add this EinsumOp to another LinearOp, returning a new LinearOp representing the sum.

static from_scalar(scalarop)

Convert a ScalarOp to an EinsumOp.

permute_for_input()

permute_for_output()

jacobian()

Return an explicit Jacobian for full-layout EinsumOp instances.

Returns:

A tensor with shape [*output_shape, *input_shape] when the operator is in full representation (is_full()). Otherwise returns NotImplemented.

Return type:

torch.Tensor

abs()

Return a LinearOp representing the element-wise absolute value of this EinsumOp.

sum_input()

Return a LinearOp that sums over the input dimensions of this EinsumOp.

sum_output()

Return a LinearOp that sums over the output dimensions of this EinsumOp.

norm_input(p=1)

Return a LinearOp that computes the norm over the input dimensions of this EinsumOp.

norm_output(p=1)

Return a LinearOp that computes the norm over the output dimensions of this EinsumOp.

property mul_conditions: EQCondition

Return a list of (input_dim, output_dim) pairs that are multiplied together in this EinsumOp.

add_conditions(target)

Return a new EinsumOp with additional (input_dim, output_dim) pairs multiplied together.

remove_conditions(target)

Return a new EinsumOp with some (input_dim, output_dim) pairs no longer multiplied together.

purify_with(other)

with_tensor(tensor)

Return a new EinsumOp with the same input/output dims but a different tensor.

__call__(x)

Apply this LinearOp to an expression, returning a Linear.

einsum_op()

Materialize this LinearOp as an explicit Jacobian tensor.

Returns:

A tensor with shape [*output_shape, *input_shape] representing the Jacobian of this LinearOp.

Return type:

EinsumOp

Notes

This may be expensive in time and memory and is mainly intended for debugging, validation, or rare paths that require explicit Jacobians.

force_jacobian()

Materialize Jacobian via batched forward/backward application.

This fallback constructs an identity basis and applies either vforward() or vbackward() depending on whether the input or output side is smaller.

Returns:

A dense Jacobian tensor with shape [*output_shape, *input_shape].

Notes

This may be expensive in time and memory and is mainly intended for debugging, validation, or rare paths that require explicit Jacobians.

jacobian_scatter(src)

Add this operator’s Jacobian contribution into an existing tensor.

Parameters:

src (torch.Tensor) – A tensor with Jacobian layout [*output_shape, *input_shape] that acts as the accumulation buffer.

Returns:

A tensor with the same shape as src containing src + jacobian(self).

Return type:

torch.Tensor

Notes

Subclasses may override this to implement structured/sparse updates without materializing the full Jacobian first.

input_shape: torch.Size

Expected input tensor shape.

output_shape: torch.Size

Computed output tensor shape.