SegmentedOperand#

class cuequivariance.SegmentedOperand#

A segmented operand is a list of segment’s shapes.

__init__(segments=None, *, ndim=None, _dims=None)#
Parameters:
classmethod empty_segments(num_segments)#

Create an operand with ndim=0

Parameters:

num_segments (int)

Return type:

SegmentedOperand

classmethod stack(operands)#

Stack a list of operands together.

Parameters:

operands (list[SegmentedOperand])

Return type:

SegmentedOperand

copy()#

Copy the operand.

Return type:

SegmentedOperand

assert_valid()#

Assert that the operand is valid.

insert_segment(index, segment)#

Insert a segment at a given index.

Parameters:
add_segment(segment)#

Add a segment to the operand.

Parameters:

segment (tuple[int, ...])

Return type:

int

property num_segments: int#

The number of segments in the operand.

property size: int#

The total size of the operand.

segment_slices()#

Return slice object for each segment.

Return type:

list[slice]

get_dimensions_dict()#

Return a dictionary of dimensions for each channel.

Return type:

dict[int, set[int]]

get_dims(i)#

Return the dimensions for a given channel.

Parameters:

i (int)

Return type:

set[int]

all_same_segment_shape()#

Check if all segments have the same shape. Returns False if there are no segments.

Return type:

bool

property segment_shape: tuple[int, ...]#

The shape of the segments if they are all the same.

property segment_size: int#

The size of the segments if they are all the same.

property slice_by_segment: _SegmentSlicer#

Return a slicer that allows slicing by segment index.

property slice_by_size: _SizeSlicer#

Return a slicer that allows slicing by flat size/offset.