SegmentedTensorProduct#

class cuequivariance.SegmentedTensorProduct#

Irreps-agnostic and dataflow-agnostic descriptor of a segmented tensor product

Parameters:

  • operands (list of operands) – The operands of the tensor product. To each operand corresponds subscripts and a list of segments.

  • paths (list of paths) – Each path contains coefficients and a list of indices. The indices are the indices of the segments of the operands.

  • coefficient_subscripts (str) – The subscripts of the coefficients.

We typically use the from_subscripts method to create a descriptor and then add segments and paths one by one.

Methods

__init__(
operands_and_subscripts=None,
coefficient_subscripts='',
*,
paths=None,
)#
Parameters:
classmethod from_subscripts(subscripts)#

Create a descriptor from a subscripts string.

Examples

>>> d = cue.SegmentedTensorProduct.from_subscripts("uv,ui,vj+ij")
>>> i0 = d.add_segment(0, (2, 3))
>>> i1 = d.add_segment(1, (2, 5))
>>> i2 = d.add_segment(2, (3, 4))
>>> d.add_path(i0, i1, i2, c=np.ones((5, 4)))
0
>>> print(d)
uv,ui,vj+ij operands=[(2, 3)],[(2, 5)],[(3, 4)] paths=[op0[0]*op1[0]*op2[0]*c c.shape=(5, 4) c.nnz=20]
Parameters:

subscripts (Subscripts)

Return type:

SegmentedTensorProduct

classmethod empty_segments(num_segments)#

Create a descriptor with a simple structure.

Examples

>>> cue.SegmentedTensorProduct.empty_segments([2, 3, 4])
,, sizes=2,3,4 num_segments=2,3,4 num_paths=0
Parameters:

num_segments (list[int])

Return type:

SegmentedTensorProduct

classmethod from_dict(data)#

Create a descriptor from a dictionary.

Parameters:

data (dict[str, Any])

Return type:

SegmentedTensorProduct

classmethod from_json(data)#

Create a descriptor from a JSON string.

Parameters:

data (str)

Return type:

SegmentedTensorProduct

classmethod from_bytes(data)#

Create a descriptor from compressed binary data.

Parameters:

data (bytes)

Return type:

SegmentedTensorProduct

classmethod from_base64(data)#

Create a descriptor from a base64 string.

Parameters:

data (str)

Return type:

SegmentedTensorProduct

property num_operands: int#

Number of operands.

property num_paths: int#

Number of paths.

property subscripts: Subscripts#

Subscripts of the tensor product.

property indices: ndarray#

Indices of the paths.

property coefficients_are_stackable: bool#

Check if the coefficients are stackable.

property stacked_coefficients: ndarray#

Stacked coefficients of the paths in a single array.

to_text(
coefficient_formatter=<function SegmentedTensorProduct.<lambda>>,
)#

Human-readable text representation of the descriptor.

Parameters:

coefficient_formatter (callable, optional) – A function to format the coefficients.

Return type:

str

Examples

>>> ((_, d),) = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1")
... ).polynomial.operations
>>> d = d.flatten_coefficient_modes()
>>> print(d.to_text())
uvw,u,v,w sizes=320,16,16,16 num_segments=5,4,4,4 num_paths=16 u=4 v=4 w=4
operand #0 subscripts=uvw
  | u: [4] * 5
  | v: [4] * 5
  | w: [4] * 5
operand #1 subscripts=u
  | u: [4] * 4
operand #2 subscripts=v
  | v: [4] * 4
operand #3 subscripts=w
  | w: [4] * 4
Flop cost: 0->1344 1->2368 2->2368 3->2368
Memory cost: 368
Path indices: 0 0 0 0, 1 0 1 1, 1 0 2 2, 1 0 3 3, 2 1 0 1, 2 2 0 2, ...
Path coefficients: [0.17...]
to_dict(extended=False)#

Dictionary representation of the descriptor.

Parameters:

extended (bool)

Return type:

dict[str, Any]

to_json(extended=False)#

JSON representation of the descriptor.

Parameters:

extended (bool)

Return type:

str

to_bytes(extended=False)#

Compressed binary representation of the descriptor.

Parameters:

extended (bool)

Return type:

bytes

to_base64(extended=False)#

Base64 representation of the descriptor.

Examples

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1")
... ).polynomial.operations[0][1]
>>> print(d.to_base64())
eJytkstuwjAQRX/F8r...lTF2zlX91/fHyvj2Z4=
Parameters:

extended (bool)

Return type:

str

get_dimensions_dict()#

Get the dimensions of the tensor product.

Return type:

dict[str, set[int]]

get_dims(m)#

Get the dimensions of a specific mode.

Examples

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+8x1"),
...     cue.Irreps("SO3", "3x0+3x1"),
...     cue.Irreps("SO3", "5x0+7x1")
... ).polynomial.operations[0][1]
>>> d.get_dims("u")
{8, 4}
>>> d.get_dims("v")
{3}
>>> d.get_dims("w")
{5, 7}
Parameters:

m (str)

Return type:

set[int]

get_path_dimensions_dict(
path,
*,
returns_sets=False,
)#

Get the dimensions of a specific path.

Parameters:

  • path (int | Path)

  • returns_sets (bool)

Return type:

dict[str, int | set[int]]

get_path_dim(path, m, *, returns_set=False)#

Get the dimension of a specific mode in a specific path.

Parameters:
Return type:

int | set[int]

segment_slice(operand, path)#

Get the slice of the segment in the given operand selected by the given path.

Parameters:

  • operand (int)

  • path (int | Path)

Return type:

slice

get_segment_shape(operand, path)#

Get the shape of the segment in the given operand selected by the given path.

Parameters:

  • operand (int)

  • path (int | Path)

Return type:

tuple[int, …]

all_segments_are_used()#

Check if all segments are used in the tensor product.

Return type:

bool

compressed_path_segment(operand)#

Starting indices of paths for the segments of the specified operand.

Note: This method requires that the paths are sorted by the specified operand.

Parameters:

operand (int) – The index of the operand for which to find path groups.

Returns:

An array of starting path indices for each segment in the specified operand.

Return type:

np.ndarray

Examples

indices[:, operand], operands[operand].num_segments -> compressed_path_segment(operand)
[0, 0, 1, 1, 1], 2 -> [0, 2, 5]
[0, 0, 1, 1, 1], 3 -> [0, 2, 5, 5]
[0, 0, 2, 2, 2], 3 -> [0, 2, 2, 5]
[1, 1], 2 -> [0, 0, 2]
operands_with_identical_segments()#

Groups of operands sharing the same segments.

Return type:

frozenset[frozenset[int]]

symmetries()#

List of permutations that leave the tensor product invariant.

Return type:

list[tuple[int, …]]

coefficients_equal_one()#

Check if all coefficients are equal to one.

Return type:

bool

flop(
operand,
batch_size=1,
algorithm='optimal',
)#

Compute the number of flops needed to compute the specified operand.

Parameters:

  • operand (int) – The operand for which to compute the flop cost.

  • batch_size (int, optional) – The batch size for the computation. Defaults to 1.

  • algorithm (str, optional) – The algorithm to use to compute the cost. Can be ‘optimal’ or ‘naive’.

Returns:

The number of flops needed to compute the specified operand.

Return type:

int

memory(batch_sizes)#

Compute the memory usage of the tensor product.

Parameters:

batch_sizes (list[int])

Return type:

int

insert_path(
path_index,
*segments,
c,
dims=None,
)#

Insert a path at a specific index.

Parameters:
Return type:

int

add_path(*segments, c, dims=None)#

Add a path to the descriptor.

Parameters:

  • segments (int | tuple[int, ...] | dict[str, int] | None) – Specifies the segments of the operands that are contracted in the path.

  • c (np.ndarray) – The coefficients of the path.

  • dims (dict[str, int], optional) – The extent of the modes.

Returns:

The index of the added path.

Return type:

int

Examples

>>> d = cue.SegmentedTensorProduct.from_subscripts("uv,ui,vj+ij")
>>> i1 = d.add_segment(1, (2, 3))
>>> i2 = d.add_segment(2, (2, 5))

We can use None to add a new segment on the fly: >>> d.add_path(None, i1, i2, c=np.ones((3, 5))) 0

The descriptor has now a new segment (2, 2) in the first operand: >>> d uv,ui,vj+ij operands=[(2, 2)],[(2, 3)],[(2, 5)] paths=[op0[0]*op1[0]*op2[0]*c c.shape=(3, 5) c.nnz=15] >>> d.add_path(0, None, None, c=np.ones((10, 10))) 1 >>> d uv,ui,vj+ij sizes=4,26,30 num_segments=1,2,2 num_paths=2 i={3, 10} j={5, 10} u=2 v=2

When the dimensions of the modes cannot be inferred, we can provide them: >>> d.add_path(None, None, None, c=np.ones((2, 2)), dims={“u”: 2, “v”: 2}) 2 >>> d uv,ui,vj+ij sizes=8,30,34 num_segments=2,3,3 num_paths=3 i={2, 3, 10} j={2, 5, 10} u=2 v=2

insert_segments(operand, sid, segments)#

Insert segments at a specific index.

Parameters:
add_segment(operand, segment)#

Add a segment to the descriptor.

Parameters:
Return type:

int

add_segments(operand, segments)#

Add segments to the descriptor.

Parameters:
canonicalize_subscripts()#

Return a new descriptor with a canonical representation of the subscripts.

Examples

>>> d = cue.SegmentedTensorProduct.from_subscripts("ab,ax,by+yx")
>>> d.canonicalize_subscripts()
uv,ui,vj+ij sizes=0,0,0 num_segments=0,0,0 num_paths=0 i= j= u= v=

This is useful to identify equivalent descriptors.

Return type:

SegmentedTensorProduct

add_or_rename_modes(
subscripts,
*,
mapping=None,
)#

Return a new descriptor with the modes renamed according to the new subscripts.

Parameters:

  • subscripts (str) – The new subscripts that contains the new names of the modes. The new subscripts can also be a superset of the old subscripts.

  • mapping (dict of str to str, optional) – The mapping between the old and new modes.

Returns:

The new descriptor with the renamed modes.

Return type:

SegmentedTensorProduct

add_or_transpose_modes(subscripts, dims=None)#

Return a new descriptor with the modes transposed according to the new subscripts.

Parameters:

  • subscripts (str) – A new subscripts that contains a permutation of the modes of the current subscripts.

  • dims (dict of str to int, optional) – The dimensions of the new modes.

Returns:

The new descriptor with the transposed modes.

Return type:

SegmentedTensorProduct

append_modes_to_all_operands(modes, dims)#

Return a new descriptor with the modes appended (to the right) to all operands.

Parameters:

  • modes (str) – The new segment modes to append to all operands.

  • dims (dict of str to int) – The dimensions of the new modes.

Returns:

The new descriptor with the appended modes.

Return type:

SegmentedTensorProduct

permute_operands(perm)#

Permute the operands of the descriptor.

Parameters:

perm (tuple[int, ...])

Return type:

SegmentedTensorProduct

move_operand(operand, new_index)#

Move an operand to a new index.

Parameters:

  • operand (int)

  • new_index (int)

Return type:

SegmentedTensorProduct

move_operand_first(operand)#

Move an operand to the first position.

Parameters:

operand (int)

Return type:

SegmentedTensorProduct

move_operand_last(operand)#

Move an operand to the last position.

Parameters:

operand (int)

Return type:

SegmentedTensorProduct

permute_segments(operand, perm)#

Permute the segments of an operand.

Parameters:
Return type:

SegmentedTensorProduct

sort_paths(operands_ordering=None)#

Sort the paths by their indices in lexicographical order.

Parameters:

operands_ordering (int or sequence of int, optional) – The order in which to sort the paths. If an integer, sort by that operand. If a sequence, sort by the first operand in the sequence and then by the second operand if equal, etc.

Returns:

The sorted descriptor.

Return type:

SegmentedTensorProduct

squeeze_modes(modes=None)#

Squeeze the descriptor by removing dimensions that are always 1.

Parameters:

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

Returns:

The squeezed descriptor.

Return type:

SegmentedTensorProduct

split_mode(mode, size)#

Split a mode into multiple modes of a given size.

Parameters:

  • mode (str) – The mode to split. All its dimensions must be divisible by the size.

  • size (int) – The size of the new modes.

Returns:

The new descriptor.

Return type:

SegmentedTensorProduct

all_same_segment_shape()#

Check if all segments have the same shape.

Return type:

bool

normalize_paths_for_operand(operand)#

Normalize the paths for an operand.

Parameters:

operand (int) – The index of the operand to normalize.

Return type:

SegmentedTensorProduct

Assuming that the input operand have unit variance, this method computes the variance of each segment in the selected operand and uniformly normalize the coefficients of the paths by the square root of the variance of the segment. This is useful to ensure that the output has unit variance.

remove_zero_paths()#

Remove paths with zero coefficients.

Return type:

SegmentedTensorProduct

fuse_paths_with_same_indices()#

Fuse paths with the same indices.

Return type:

SegmentedTensorProduct

consolidate_paths()#

Consolidate the paths by merging duplicates and removing zeros.

Return type:

SegmentedTensorProduct

sort_indices_for_identical_operands(operands)#

Reduce the number of paths by sorting the indices for identical operands.

Parameters:

operands (Sequence[int])

Return type:

SegmentedTensorProduct

symmetrize_operands(operands, force=False)#

Symmetrize the specified operands permuting the indices.

Parameters:
Return type:

SegmentedTensorProduct

remove_empty_segments()#

Remove empty segments.

Return type:

SegmentedTensorProduct

flatten_modes(
modes,
*,
skip_zeros=True,
force=False,
)#

Remove the specified modes by subdividing segments and paths.

Parameters:

  • modes (Sequence of str) – The modes to remove, they must precede the modes to keep in each operand.

  • skip_zeros (bool, optional) – Whether to skip paths with zero coefficients. Default is True.

  • force (bool, optional) – Whether to force the flattening by flattening extra necessary modes. Default is False.

Return type:

SegmentedTensorProduct

flatten_coefficient_modes(
*,
skip_zeros=True,
force=False,
)#

Flatten the coefficients of the descriptor. Create new segments and paths to flatten the coefficients.

Parameters:

  • skip_zeros (bool, optional) – Whether to skip paths with zero coefficients. Default is True.

  • force (bool, optional) – Whether to force the flattening by flattening extra necessary modes. Default is False.

Return type:

SegmentedTensorProduct

consolidate_modes(modes=None)#

Consolidate the descriptor by merging modes together.

Parameters:

modes (str, optional) – The modes to consolidate. If None, consolidate all modes that are repeated.

Return type:

SegmentedTensorProduct

round_coefficients_to_rational(max_denominator)#

Round the coefficients to the nearest p / q number with a given maximum denominator.

Parameters:

max_denominator (int) – The maximum denominator, q < max_denominator.

Return type:

SegmentedTensorProduct

round_coefficients_to_sqrt_rational(
max_denominator,
)#

Round the coefficients to the nearest sqrt(p / q) number with a given maximum denominator.

Parameters:

max_denominator (int) – The maximum denominator, q < max_denominator.

Return type:

SegmentedTensorProduct

modify_coefficients(f)#

Modify the coefficients of the descriptor.

Parameters:

f (callable) – The function to apply to the coefficients.

Return type:

SegmentedTensorProduct

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.