SegmentedPolynomial#

class cuequivariance.SegmentedPolynomial#

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.

__init__(inputs, outputs, operations)#
Parameters:
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)#

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, stacked)#

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, outputs, operations)#

Stack segmented tensor products together.

Parameters:
Returns:

The stacked polynomial.

Return type:

cue.SegmentedPolynomial

classmethod concatenate(inputs, outputs, polys)#

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()#

Check if all operands have the same segment shape.

Returns:

True if all operands have the same segment shape.

Return type:

bool

used_inputs()#

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()#

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()#

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=1)#

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)#

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

permute_inputs(permutation)#

Permute the input operands of the polynomial.

Parameters:

permutation (list of int) – The permutation to apply to the inputs.

Returns:

A new polynomial with permuted inputs.

Return type:

cue.SegmentedPolynomial

permute_outputs(permutation)#

Permute the output operands of the polynomial.

Parameters:

permutation (list of int) – The permutation to apply to the outputs.

Returns:

A new polynomial with permuted outputs.

Return type:

cue.SegmentedPolynomial

apply_fn(f)#

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()#

Fuse segmented tensor products with identical operations and operands.

Returns:

Polynomial with fused tensor products.

Return type:

cue.SegmentedPolynomial

consolidate()#

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

Returns:

Consolidated polynomial.

Return type:

cue.SegmentedPolynomial

flatten_modes(modes)#

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()#

Canonicalize the subscripts of the segmented tensor products.

Returns:

Polynomial with canonicalized subscripts.

Return type:

cue.SegmentedPolynomial

squeeze_modes(modes=None)#

Squeeze modes that are always 1 in all operations.

Parameters:

modes (str, optional) – The modes to squeeze. If None, squeeze all modes that are always 1.

Returns:

Polynomial with squeezed modes.

Return type:

cue.SegmentedPolynomial

Note

When modes is None, all operands with 1-dimensions are squeezed, including unused operands. When modes is specified, unused operands (not linked to any operation) are not squeezed because they don’t have mode information to determine which dimensions correspond to which modes.

split_mode(mode, size)#

Split specified mode in the polynomial.

Parameters:

  • mode (str) – Mode to split.

  • size (int) – Size to split the mode into.

Returns:

Polynomial with split mode.

Return type:

cue.SegmentedPolynomial

flatten_coefficient_modes()#

Flatten the coefficient modes of the segmented tensor products.

Returns:

Polynomial with flattened coefficient modes.

Return type:

cue.SegmentedPolynomial

symmetrize_for_identical_operands()#

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()#

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

split_operand_by_segment(
operand_id,
segment_splits,
)#

Split an operand into multiple operands based on segment boundaries.

Parameters:

  • operand_id (int) – Index of the operand to split.

  • segment_splits (list of int) – List of segment indices where to split the operand. Must start with 0 and end with the total number of segments.

Returns:

Polynomial with the specified operand split.

Return type:

cue.SegmentedPolynomial

split_operand_by_size(operand_id, offsets)#

Split an operand into multiple operands based on specified offsets.

Parameters:

  • operand_id (int) – Index of the operand to split.

  • offsets (list of int) – List of offsets to split the operand at.

Returns:

Polynomial with the specified operand split.

Return type:

cue.SegmentedPolynomial

filter_keep_operands(keep)#

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_unused_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)#

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_unused_operands()#

Remove all unused operands from the polynomial.

Returns:

Polynomial with unused operands removed.

Return type:

cue.SegmentedPolynomial

compute_only(keep)#

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)#

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, has_cotangent)#

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, has_cotangent)#

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