SegmentedPolynomial#

class cuequivariance.SegmentedPolynomial(
inputs: Sequence[SegmentedOperand],
outputs: Sequence[SegmentedOperand],
operations: Sequence[tuple[Operation | Sequence[int], SegmentedTensorProduct]],
)#

A polynomial representation using segmented tensor products.

This class represents a polynomial using a collection of segmented tensor products, where each product is associated with an operation that specifies how inputs are combined. The polynomial maps a set of input tensors to output tensors through these tensor products.

Parameters:

  • inputs (tuple of SegmentedOperand) – Input operands.

  • outputs (tuple of SegmentedOperand) – Output operands.

  • operations (list of tuple of Operation and SegmentedTensorProduct) – List of operation and tensor product pairs that define the polynomial transformation.

property operands: tuple[SegmentedOperand, ...]#

Get all operands (inputs and outputs) of the polynomial.

Returns:

Tuple of all operands.

Return type:

tuple of cue.SegmentedOperand

property num_inputs: int#

Get the number of input operands.

Returns:

Number of input operands.

Return type:

int

property num_outputs: int#

Get the number of output operands.

Returns:

Number of output operands.

Return type:

int

property num_operands: int#

Get the total number of operands (inputs + outputs).

Returns:

Total number of operands.

Return type:

int

classmethod eval_last_operand(
stp: SegmentedTensorProduct,
)#

Create a polynomial that evaluates the last operand of a segmented tensor product.

Parameters:

stp (cue.SegmentedTensorProduct) – The tensor product to evaluate.

Returns:

A polynomial evaluating the last operand.

Return type:

cue.SegmentedPolynomial

classmethod stack(
polys: list[SegmentedPolynomial],
stacked: list[bool],
) → SegmentedPolynomial#

Stack segmented polynomials together.

This method combines multiple polynomials by stacking their operands where indicated by the stacked parameter. Non-stacked operands must be identical across all polynomials.

Parameters:
Returns:

The stacked polynomial.

Return type:

cue.SegmentedPolynomial

Example

>>> p1 = cue.descriptors.spherical_harmonics(cue.SO3(1), [1, 2]).polynomial
>>> p2 = cue.descriptors.spherical_harmonics(cue.SO3(1), [2, 3]).polynomial
>>> cue.SegmentedPolynomial.stack([p1, p2], [False, True])
╭ a=[3:3⨯()] -> B=[20:20⨯()]
│  []·a[]➜B[] ───────── num_paths=3
│  []·a[]·a[]➜B[] ───── num_paths=22
╰─ []·a[]·a[]·a[]➜B[] ─ num_paths=41

Note how the STPs of degree 2 has been automatically fused into a single STP.

classmethod stack_tensor_products(
inputs: Sequence[SegmentedOperand | None],
outputs: Sequence[SegmentedOperand | None],
operations: Sequence[tuple[Operation | Sequence[int], SegmentedTensorProduct]],
) → SegmentedPolynomial#

Stack segmented tensor products together.

Parameters:
Returns:

The stacked polynomial.

Return type:

cue.SegmentedPolynomial

classmethod concatenate(
inputs: Sequence[SegmentedOperand],
outputs: Sequence[SegmentedOperand],
polys: list[tuple[SegmentedPolynomial, Sequence[int | None]]],
) → SegmentedPolynomial#

Concatenate segmented polynomials.

Parameters:

  • inputs (list of cue.SegmentedOperand) – Input operands for the concatenated polynomial.

  • outputs (list of cue.SegmentedOperand) – Output operands for the concatenated polynomial.

  • polys (list of tuple of cue.SegmentedPolynomial and list of int | None) – List of tuples containing (polynomial, mapping), where mapping[i] is the operand index in the polynomial that corresponds to the i-th operand in the concatenated polynomial. If mapping[i] is None, the i-th operand in the concatenated polynomial is not used in the polynomial.

Returns:

The concatenated polynomial.

Return type:

cue.SegmentedPolynomial

Example

>>> p1 = cue.descriptors.spherical_harmonics(cue.SO3(1), [1, 2]).polynomial
>>> p2 = cue.descriptors.spherical_harmonics(cue.SO3(1), [2, 3]).polynomial
>>> [vec, sh1] = p1.operands
>>> [_, sh2] = p2.operands
>>> cue.SegmentedPolynomial.concatenate(
...     [vec],
...     [sh1, sh2],
...     [(p1, [0, 1, None]), (p2, [0, None, 1])],
... )
╭ a=[3:3⨯()] -> B=[8:8⨯()] C=[12:12⨯()]
│  []·a[]➜B[] ───────── num_paths=3
│  []·a[]·a[]➜B[] ───── num_paths=11
│  []·a[]·a[]➜C[] ───── num_paths=11
╰─ []·a[]·a[]·a[]➜C[] ─ num_paths=41
all_same_segment_shape() → bool#

Check if all operands have the same segment shape.

Returns:

True if all operands have the same segment shape.

Return type:

bool

used_inputs() → list[bool]#

Get list of boolean values indicating which inputs are used in the polynomial.

Returns:

List where True indicates the input is used.

Return type:

list of bool

used_outputs() → list[bool]#

Get list of boolean values indicating which outputs are used in the polynomial.

Returns:

List where True indicates the output is used.

Return type:

list of bool

used_operands() → list[bool]#

Get list of boolean values indicating which operands are used in the polynomial.

Returns:

List where True indicates the operand is used.

Return type:

list of bool

flop(batch_size: int = 1) → int#

Compute the number of floating point operations in the polynomial.

Parameters:

batch_size (int, optional) – Batch size for computation. Defaults to 1.

Returns:

Number of floating point operations.

Return type:

int

memory(batch_sizes: list[int]) → int#

Compute the memory usage of the polynomial.

Parameters:

batch_sizes (list of int) – List of batch sizes for each operand. Each operand can have its own batch size, allowing for different batch dimensions per tensor.

Returns:

Memory usage in number of elements.

Return type:

int

apply_fn(
f: Callable[[Operation, SegmentedTensorProduct], tuple[Operation, SegmentedTensorProduct] | None],
) → SegmentedPolynomial#

Apply a function to each tensor product in the polynomial.

Parameters:

f (Callable) – Function to apply to each operation and tensor product pair.

Returns:

New polynomial with transformed tensor products.

Return type:

cue.SegmentedPolynomial

fuse_stps() → SegmentedPolynomial#

Fuse segmented tensor products with identical operations and operands.

Returns:

Polynomial with fused tensor products.

Return type:

cue.SegmentedPolynomial

consolidate() → SegmentedPolynomial#

Consolidate the segmented tensor products by removing empty segments and squeezing modes.

Returns:

Consolidated polynomial.

Return type:

cue.SegmentedPolynomial

flatten_modes(
modes: list[str],
) → SegmentedPolynomial#

Flatten specified modes in the polynomial.

Parameters:

modes (list of str) – List of mode names to flatten.

Returns:

Polynomial with flattened modes.

Return type:

cue.SegmentedPolynomial

canonicalize_subscripts() → SegmentedPolynomial#

Canonicalize the subscripts of the segmented tensor products.

Returns:

Polynomial with canonicalized subscripts.

Return type:

cue.SegmentedPolynomial

squeeze_modes(
modes: str | None = None,
) → SegmentedPolynomial#

Squeeze specified modes in the polynomial.

Parameters:

modes (str | None, optional) – Modes to squeeze. If None, squeezes all modes.

Returns:

Polynomial with squeezed modes.

Return type:

cue.SegmentedPolynomial

flatten_coefficient_modes() → SegmentedPolynomial#

Flatten the coefficient modes of the segmented tensor products.

Returns:

Polynomial with flattened coefficient modes.

Return type:

cue.SegmentedPolynomial

symmetrize_for_identical_operands() → SegmentedPolynomial#

Symmetrize the paths of the segmented tensor products for identical operands.

This operation increases the number of paths in the segmented tensor products.

Returns:

Polynomial with symmetrized paths.

Return type:

cue.SegmentedPolynomial

unsymmetrize_for_identical_operands() → SegmentedPolynomial#

Unsymmetrize the paths of the segmented tensor products for identical operands.

This operation decreases the number of paths in the segmented tensor products.

Returns:

Polynomial with unsymmetrized paths.

Return type:

cue.SegmentedPolynomial

filter_keep_operands(
keep: list[bool],
) → SegmentedPolynomial#

Select which operands to keep in the polynomial.

Use this method when you want to compute only a subset of the polynomial outputs and have control over which inputs to keep. For keeping all inputs (even if not used), use filter_keep_outputs. For automatically removing unused operands, use filter_drop_unsued_operands.

Parameters:

keep (list of bool) – List indicating which operands to keep.

Returns:

Polynomial with selected operands.

Return type:

cue.SegmentedPolynomial

filter_keep_outputs(
keep: list[bool],
) → SegmentedPolynomial#

Select which outputs to keep in the polynomial.

Parameters:

keep (list[bool]) – List indicating which outputs to keep.

Returns:

Polynomial with selected outputs.

Return type:

cue.SegmentedPolynomial

filter_drop_unsued_operands() → SegmentedPolynomial#

Remove all unused operands from the polynomial.

Returns:

Polynomial with unused operands removed.

Return type:

cue.SegmentedPolynomial

compute_only(
keep: list[bool],
) → SegmentedPolynomial#

Create a polynomial that only computes selected outputs.

The new polynomial will keep the same operands as the original one, but will only compute the selected outputs.

Parameters:

keep (list of bool) – List indicating which outputs to compute.

Returns:

Polynomial computing only selected outputs.

Return type:

cue.SegmentedPolynomial

jvp(
has_tangent: list[bool],
) → tuple[SegmentedPolynomial, Callable[[tuple[list[Any], list[Any]]], tuple[list[Any], list[Any]]]]#

Compute the Jacobian-vector product of the polynomial.

Parameters:

has_tangent (list of bool) – List indicating which inputs have tangents.

Returns:

The JVP polynomial and a mapping function for inputs/outputs.

Return type:

tuple of cue.SegmentedPolynomial and Callable

transpose(
is_undefined_primal: list[bool],
has_cotangent: list[bool],
) → tuple[SegmentedPolynomial, Callable[[tuple[list[Any], list[Any]]], tuple[list[Any], list[Any]]]]#

Transpose the polynomial for reverse-mode automatic differentiation.

Parameters:

  • is_undefined_primal (list of bool) – List indicating which inputs have undefined primals.

  • has_cotangent (list of bool) – List indicating which outputs have cotangents.

Returns:

The transposed polynomial and a mapping function for inputs/outputs.

Return type:

tuple of cue.SegmentedPolynomial and Callable

backward(
requires_gradient: list[bool],
has_cotangent: list[bool],
) → tuple[SegmentedPolynomial, Callable[[tuple[list[Any], list[Any]]], tuple[list[Any], list[Any]]]]#

Compute the backward pass of the polynomial for gradient computation.

Parameters:

  • requires_gradient (list of bool) – List indicating which inputs require gradients.

  • has_cotangent (list of bool) – List indicating which outputs have cotangents.

Returns:

The backward polynomial and a mapping function for inputs/outputs.

Return type:

tuple of cue.SegmentedPolynomial and Callable