EquivariantPolynomial#

class cuequivariance.EquivariantPolynomial( inputs: list[Rep], outputs: list[Rep], polynomial: SegmentedPolynomial, )#

A polynomial representation with equivariance constraints.

This class extends cue.SegmentedPolynomial by incorporating information about the group representations of each input and output tensor. It ensures that operations performed by the polynomial respect the equivariance constraints defined by these representations, making it suitable for building equivariant neural networks.

Parameters:

inputs (tuple of cue.Rep) – Group representations for input tensors.
outputs (tuple of cue.Rep) – Group representations for output tensors.
polynomial (cue.SegmentedPolynomial) – The underlying polynomial transformation.

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

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

Returns:: Tuple of all operands.
Return type:: tuple of cue.Rep

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 stack( polys: list[EquivariantPolynomial], stacked: list[bool], ) → EquivariantPolynomial#

Stack equivariant 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:

polys (list of cue.EquivariantPolynomial) – List of polynomials to stack.
stacked (list of bool) – List indicating which operands should be stacked.

Returns:

The stacked polynomial.

Return type:

cue.EquivariantPolynomial

Example

>>> p1 = cue.descriptors.spherical_harmonics(cue.SO3(1), [1, 2])
>>> p2 = cue.descriptors.spherical_harmonics(cue.SO3(1), [2, 3])
>>> cue.EquivariantPolynomial.stack([p1, p2], [False, True])
╭ a=1 -> B=1+2+2+3
│  []·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.

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], ) → EquivariantPolynomial#

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.EquivariantPolynomial

fuse_stps() → EquivariantPolynomial#

Fuse segmented tensor products with identical operations and operands.

Returns:: Polynomial with fused tensor products.
Return type:: cue.EquivariantPolynomial

consolidate() → EquivariantPolynomial#

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

Returns:: Consolidated polynomial.
Return type:: cue.EquivariantPolynomial

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

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.EquivariantPolynomial

canonicalize_subscripts() → EquivariantPolynomial#

Canonicalize the subscripts of the segmented tensor products.

Returns:: Polynomial with canonicalized subscripts.
Return type:: cue.EquivariantPolynomial

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

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.EquivariantPolynomial

flatten_coefficient_modes() → EquivariantPolynomial#

Flatten the coefficient modes of the segmented tensor products.

Returns:: Polynomial with flattened coefficient modes.
Return type:: cue.EquivariantPolynomial

symmetrize_for_identical_operands() → EquivariantPolynomial#

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.EquivariantPolynomial

unsymmetrize_for_identical_operands() → EquivariantPolynomial#

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.EquivariantPolynomial

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

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.EquivariantPolynomial

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

Select which outputs to keep in the polynomial.

Parameters:: keep (list of bool) – List indicating which outputs to keep.
Returns:: Polynomial with selected outputs.
Return type:: cue.EquivariantPolynomial

filter_drop_unsued_operands() → EquivariantPolynomial#

Remove all unused operands from the polynomial.

Returns:: Polynomial with unused operands removed.
Return type:: cue.EquivariantPolynomial

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

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.EquivariantPolynomial

jvp( has_tangent: list[bool], ) → tuple[EquivariantPolynomial, 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.EquivariantPolynomial and Callable

transpose( is_undefined_primal: list[bool], has_cotangent: list[bool], ) → tuple[EquivariantPolynomial, 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.EquivariantPolynomial and Callable

backward( requires_gradient: list[bool], has_cotangent: list[bool], ) → tuple[EquivariantPolynomial, 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.EquivariantPolynomial and Callable