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.
backward	Apply the transposed linear function to an input 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.
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_op	Materialize this LinearOp as an explicit Jacobian tensor.
jacobian_scatter	Accumulate this operator's Jacobian into src efficiently.
permute_for_input
permute_for_output
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.
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.

__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, 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.

__add__(other)

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

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.

jacobian_scatter(src)

Accumulate this operator’s Jacobian into src efficiently.

Parameters:

src (torch.Tensor) – Accumulator tensor with shape [*output_shape, *input_shape].

Returns:

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

Return type:

torch.Tensor

Notes

For full-layout operators, this reduces to direct addition. For compressed/masked layouts, this method performs a structured diagonal-scatter update that avoids full Jacobian expansion in the common case.

__call__(x)

Apply this LinearOp to an expression, returning a Linear.

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_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:

torch.Tensor

Notes

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

input_shape: torch.Size

Expected input tensor shape.

output_shape: torch.Size

Computed output tensor shape.