NVIDIA cuEquivariance

https://docs.nvidia.com/cuda/cuequivariance/
Admin
cuEquivariance is an NVIDIA Python library for building high-performance geometric neural networks...
Tokens:108,813
Snippets:873
Trust Score:10
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install cuEquivariance with pip

Source: https://docs.nvidia.com/cuda/cuequivariance/index.html

Install cuEquivariance using pip, choosing the frontend (JAX, PyTorch, or core) and CUDA kernel versions as needed. Ensure compatibility with your PyTorch or JAX versions.

```bash
pip install cuequivariance-jax
pip install cuequivariance-torch
pip install cuequivariance  # Installs only the core non-ML components

# CUDA kernels
pip install cuequivariance-ops-torch-cu12  # or -cu13
pip install cuequivariance-ops-jax-cu12    # or -cu13
```

--------------------------------

### Linear Layer Example

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.SegmentedPolynomial.html

Shows how to use SegmentedPolynomial to create and apply a linear transformation with random weights and input. The 'naive' method is used for demonstration.

```python
# Create a linear transformation
input_irreps = cue.Irreps(cue.O3, "5x0e + 3x1o")
output_irreps = cue.Irreps(cue.O3, "4x0e + 2x1o")
poly = cue.descriptors.linear(input_irreps, output_irreps).polynomial

# Create the module
linear = SegmentedPolynomial(poly, method="naive")

# Create random weights and input
weights = torch.randn(1, poly.inputs[0].size)
x = torch.randn(10, poly.inputs[1].size)

# Forward pass
result = linear([weights, x])
print(result[0].shape)
```

--------------------------------

### Basic Segmented Polynomial Execution in JAX

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/jax/poly.html

Demonstrates the fundamental steps to create and execute a Segmented Polynomial in JAX. Ensure cuequivariance and cuequivariance_jax are installed.

```python
import jax
import jax.numpy as jnp
import cuequivariance as cue
import cuequivariance_jax as cuex

# 1. Create a polynomial (e.g. x^2)
stp = cue.SegmentedTensorProduct.from_subscripts("i,j,ij")
stp.add_segment(0, (3,))
stp.add_segment(1, (3,))
stp.add_segment(2, (3, 3))
stp.add_path(0, 0, 0, c=1.0)

x = stp.operands[0]
sp = cue.SegmentedPolynomial(
    inputs=[x],
    outputs=[stp.operands[2]],
    operations=[(cue.Operation([0, 0, 1]), stp)]
)

# 2. Define output shapes/dtypes
# We need to tell JAX what the output will look like
# Use -1 for the last dimension to infer from the polynomial descriptor
output_structs = [jax.ShapeDtypeStruct((-1,), jnp.float32)]

# 3. Execute
input_arr = jnp.ones((3,))
[output_arr] = cuex.segmented_polynomial(
    sp,
    [input_arr],
    output_structs,
    method="naive"
)

print(f"Output shape: {output_arr.shape}")
```

```text
Output shape: (9,)
```

--------------------------------

### Wiring a Square Function with Segmented Polynomial

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/poly.html

This example demonstrates how to create a Segmented Polynomial (SP) to compute y=x⊗x. It shows how to define the math (STP) and then wire a single global input to both operands of the STP.

```python
import cuequivariance as cue
import numpy as np
from cuequivariance.segmented_polynomials import visualize_polynomial

# 1. Create the STP (The Math)
# A simple contraction: A_i * B_j -> C_ij
stp = cue.SegmentedTensorProduct.from_subscripts("i,j,ij")
stp.add_segment(0, (3,))  # Left operand: size 3
stp.add_segment(1, (3,))  # Right operand: size 3
stp.add_segment(2, (3, 3))  # Output segment: shape (3, 3) for modes i,j
stp.add_path(0, 0, 0, c=1.0)  # Scalar coefficient

# 2. Create the SP (The Wiring)
# Global Inputs: [x] (We define x using the shape from the STP)
x = stp.operands[0]

# The Operation defines the wiring: [0, 0, 1]
#   - STP Operand 0 gets Global Input 0 (x)
#   - STP Operand 1 gets Global Input 0 (x) -- REUSE!
#   - STP Operand 2 becomes the Output

sp = cue.SegmentedPolynomial(
    inputs=[x],
    outputs=[stp.operands[2]],
    operations=[(cue.Operation([0, 0, 1]), stp)]
)

print(sp)
```

--------------------------------

### Handle No Solution for Clebsch-Gordan Coefficients

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.clebsch_gordan.html

This example demonstrates the output when no Clebsch-Gordan coefficients exist for the given representations. The shape indicates zero solutions.

```python
C = clebsch_gordan(cue.SO3(1), cue.SO3(1), cue.SO3(3))
C.shape
>>> (0, 3, 3, 7)
```

--------------------------------

### cuEquivariance Channelwise Tensor Product Setup

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html

Sets up the channelwise tensor product operation using cuEquivariance, defining the irreps and creating a SegmentedPolynomial module. This approach integrates gather and scatter operations for improved performance.

```python
# Cue version of the irreps
irreps_in1 = cue.Irreps("O3", irreps_node_input)
irreps_in2 = cue.Irreps("O3", irreps_edge_attr)
irreps_out = cue.Irreps("O3", target_irreps)
# Defining the operation
e = cue.descriptors.channelwise_tensor_product(
    irreps_in1, irreps_in2, irreps_out
)
# The TP itself:
cue_tp = cuet.SegmentedPolynomial(
    e.polynomial,
    method="uniform_1d"
).to(device)
```

--------------------------------

### Indexed Operations Example

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.SegmentedPolynomial.html

Illustrates the use of indexed operations with SegmentedPolynomial, allowing different weight sets to be used for different inputs. This example assumes 'linear' and 'poly' are already defined from the previous snippet.

```python
# Create indexed weights for different elements
weights = torch.randn(3, poly.inputs[0].size)  # 3 different weight sets
x = torch.randn(5, poly.inputs[1].size)        # 5 input vectors

# Index tensor specifying which weights to use for each input
weight_indices = torch.tensor([0, 1, 0, 2, 1])  # Use weights 0,1,0,2,1

result = linear([weights, x],
               input_indices={0: weight_indices})
print(result[0].shape)
```

--------------------------------

### Prepare Graph Data for Indexing

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/jax/ir_dict.html

Initializes graph data including weights, node features, spherical harmonics, and sender/receiver indices. This setup is for operations involving graph structures.

```python
num_edges, num_nodes = 100, 30

w = jax.random.normal(jax.random.key(0), (num_edges, poly.inputs[0].size))
w = rearrange(w, "e (s m) -> e s m", s=poly.inputs[0].num_segments)

node_feats = {
    cue.SO3(0): jax.random.normal(jax.random.key(1), (num_nodes, 1, 32)),
    cue.SO3(1): jax.random.normal(jax.random.key(2), (num_nodes, 3, 32)),
}

sph = {
    cue.SO3(0): jax.random.normal(jax.random.key(3), (num_edges, 1)),
    cue.SO3(1): jax.random.normal(jax.random.key(4), (num_edges, 3)),
}

senders = jax.random.randint(jax.random.key(5), (num_edges,), 0, num_nodes)
receivers = jax.random.randint(jax.random.key(6), (num_edges,), 0, num_nodes)
```

--------------------------------

### SymmetricContraction Initialization

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.SymmetricContraction.html

Initializes the SymmetricContraction layer with specified input and output irreps, contraction degree, and other optional parameters. The example demonstrates setting up the layer with common parameters like layout, dtype, and device.

```APIDOC
## class cuequivariance_torch.SymmetricContraction

### Description
Accelerated implementation of the symmetric contraction operation.

### Parameters
- **irreps_in** (_Irreps_) - The input irreps. All multiplicities (mul) within the irreps must be identical.
- **irreps_out** (_Irreps_) - The output irreps. All multiplicities must be the same.
- **contraction_degree** (_int_) - The degree of the symmetric contraction.
- **num_elements** (_int_) - The number of elements for the weight tensor.
- **layout** (_IrrepsLayout_, optional) - The layout of the input and output irreps. Defaults to a default layout.
- **layout_in** (_IrrepsLayout_, optional) - The layout of the input irreducible representations.
- **layout_out** (_IrrepsLayout_, optional) - The layout of the output irreducible representations.
- **device** (_torch.device_, optional) - The device to use for the operation.
- **dtype** (_torch.dtype_, optional) - The dtype to use for the operation weights. Defaults to `torch.float32`.
- **math_dtype** (_str_ or _torch.dtype_, optional) - The dtype to use for the math operations.
- **original_mace** (_bool_, optional) - Whether to use the original MACE implementation. Defaults to False.
- **method** (_str_, optional) - The method to use for the operation. Defaults to 'uniform_1d' or 'naive'.
- **use_fallback** (_bool_, optional) - Deprecated. Maps to `method` parameter.

### Example
```python
import torch
import cuequivariance_torch as cuet

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
irreps_in = cuet.Irreps("O3", "32x0e + 32x1o")
irreps_out = cuet.Irreps("O3", "32x0e")

layer = cuet.SymmetricContraction(
    irreps_in=irreps_in,
    irreps_out=irreps_out,
    contraction_degree=3,
    num_elements=5,
    layout=cuet.ir_mul, 
    dtype=torch.float32,
    device=device
)
```
```

--------------------------------

### Triangle Multiplicative Update Example

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.triangle_multiplicative_update.html

Demonstrates how to use triangle_multiplicative_update with input tensors and a mask. It shows the shape of the output and the gradients after a backward pass. Ensure CUDA is available before running.

```python
import torch
from cuequivariance_torch import triangle_multiplicative_update
if torch.cuda.is_available():
    device = torch.device("cuda")
    batch_size, seq_len, hidden_dim = 1, 128, 128
    # Create input tensor
    x = torch.randn(batch_size, seq_len, seq_len, hidden_dim, requires_grad=True, device=device)
    # Create mask (1 for valid positions, 0 for masked)
    mask = torch.ones(batch_size, seq_len, seq_len, device=device)
    # Perform triangular multiplication
    output = triangle_multiplicative_update(
        x=x,
        direction="outgoing",  # or "incoming"
        mask=mask,
    )
    print(output.shape)  # torch.Size([1, 128, 128, 128])
    # Create gradient tensor and perform backward pass
    grad_out = torch.randn_like(output)
    output.backward(grad_out)
    # Access gradients
    print(x.grad.shape)  # torch.Size([1, 128, 128, 128])
```

--------------------------------

### Initialize SymmetricContraction with Original MACE Implementation

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.SymmetricContraction.html

Instantiate the SymmetricContraction layer, optionally using the original MACE implementation by setting `original_mace=True`. This example shows the new function definition compared to an older one.

```python
feats_irreps = cue.Irreps("O3", "32x0e + 32x1o + 32x2e")
target_irreps = cue.Irreps("O3", "32x0e + 32x1o")
# OLD FUNCTION DEFINITION:
# symmetric_contractions_old = SymmetricContraction(
#     irreps_in=feats_irreps,
#     irreps_out=target_irreps,
#     correlation=3,
#     num_elements=10,
# )
# NEW FUNCTION DEFINITION:
symmetric_contractions_new = cuet.SymmetricContraction(
    irreps_in=feats_irreps,
    irreps_out=target_irreps,
    contraction_degree=3,
    num_elements=10,
```

--------------------------------

### Triangle Multiplicative Update with Default Dimensions

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_jax.triangle_multiplicative_update.html

Example demonstrating the use of triangle_multiplicative_update with default dimension preservation. Input and output dimensions remain the same. Ensure all necessary weights and biases are provided or initialized.

```python
import jax
import jax.numpy as jnp
from cuequivariance_jax import triangle_multiplicative_update

# Create input tensor with arbitrary batch dimensions
key = jax.random.key(0)
key, subkey = jax.random.split(key)
batch_dim1, batch_dim2, seq_len, D_in = 2, 3, 128, 128
x = jax.random.normal(subkey, (batch_dim1, batch_dim2, seq_len, seq_len, D_in), dtype=jnp.float32)

# Create mask (1 for valid positions, 0 for masked)
mask = jnp.ones((batch_dim1, batch_dim2, seq_len, seq_len))

# Create weight parameters (in practice, these would be learned)
norm_in_weight = jnp.ones(D_in)
norm_in_bias = jnp.zeros(D_in)

# Optional bias parameters for projection and gating layers
p_in_bias = jnp.zeros(2 * D_in)  # Optional input projection bias
g_in_bias = jnp.zeros(2 * D_in)  # Optional input gating bias
p_out_bias = jnp.zeros(D_in)     # Optional output projection bias (would be D_out if dimension changes)
g_out_bias = jnp.zeros(D_in)     # Optional output gating bias (would be D_out if dimension changes)

# Initialize other weights using the key
key, subkey = jax.random.split(key)

# Perform triangular multiplication
output = triangle_multiplicative_update(
    x=x,
    direction="outgoing",  # or "incoming"
    key=subkey,  # Only needed if some weights are None
    mask=mask,
    norm_in_weight=norm_in_weight,
    norm_in_bias=norm_in_bias,
    p_in_bias=p_in_bias,  # Can be None to skip bias
    g_in_bias=g_in_bias,  # Can be None to skip bias
    p_out_bias=p_out_bias,  # Can be None to skip bias
    g_out_bias=g_out_bias,  # Can be None to skip bias
    # ... pass other weights or let them initialize ...
)

print(output.shape)
```

--------------------------------

### SegmentedPolynomial with Input and Output Indexing

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/poly.html

Demonstrates advanced indexing for routing data flexibly, such as applying different weights to different examples or accumulating results into specific output bins. Input indices map operand indices to index tensors, while output indices map results to bins.

```python
# Example: Select from 3 weight sets for 10 input examples
num_weights, num_examples, dim = 3, 10, 32
weights = torch.randn(num_weights, dim)
inputs = torch.randn(num_examples, dim)

# Index tensor: Assign weight set 0 to first 3 examples, etc.
w_idx = torch.tensor([0, 0, 0, 1, 1, 1, 1, 2, 2, 2])

# Execute
# input_indices maps Input 0 (weights) to w_idx
[z] = model(
    [weights, inputs],
    input_indices={0: w_idx}
)

print(f"Output shape: {z.shape}")

# Output Indexing: Accumulate results into specific bins
# We want 5 output bins. We map the 10 results to bins [0, 0, 1, 1, ...]
out_idx = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4])
output_shape = torch.empty(5, dim)

[z_accum] = model(
    [weights, inputs],
    input_indices={0: w_idx},
    output_indices={0: out_idx},
    output_shapes={0: output_shape}
)

print(f"Accumulated Output shape: {z_accum.shape}")
```

```text
Output shape: torch.Size([10, 32])
Accumulated Output shape: torch.Size([5, 32])
```

--------------------------------

### Triangle Multiplicative Update with Dimension Change

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_jax.triangle_multiplicative_update.html

Example demonstrating triangle_multiplicative_update when changing the output dimension. This requires providing specific weights (g_out_weight, p_out_weight) with the desired output dimension. The input dimension is preserved in the intermediate calculations.

```python
# Example with dimension change: input 128 -> output 256
g_out_weight_256 = jax.random.normal(jax.random.key(1), (256, 128))
p_out_weight_256 = jax.random.normal(jax.random.key(2), (256, 128))

# Assuming x, key, and subkey2 are defined as in the previous example
# key, subkey2 = jax.random.split(key)

output_256 = triangle_multiplicative_update(
    x=x,
    direction="outgoing",
    key=subkey2,  # Key needed for other weight initialization
    g_out_weight=g_out_weight_256,
    p_out_weight=p_out_weight_256,
)

print(output_256.shape)
```

--------------------------------

### Attention with Caching of Projected Z

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.attention_pair_bias.html

Example demonstrating how to use caching for the projected z tensor, recommended for inference when z remains constant. It checks for existing cache, calls the function, and updates the cache if necessary.

```python
# Check cache and determine if z is already projected
if model_cache is not None and "proj_z" in model_cache:
    z = model_cache["proj_z"]
    is_cached_z = True
else:
    is_cached_z = False

# Call attention_pair_bias
o, proj_z = attention_pair_bias(
    s=s, q=q, k=k, v=v, z=z, mask=mask,
    num_heads=num_heads,
    w_proj_z=w_proj_z if not is_cached_z else None,
    w_proj_g=w_proj_g,
    w_proj_o=w_proj_o,
    w_ln_z=w_ln_z if not is_cached_z else None,
    b_ln_z=b_ln_z if not is_cached_z else None,
    return_z_proj=True,
    is_cached_z_proj=is_cached_z,
)

# Cache proj_z for next call
if model_cache is not None and "proj_z" not in model_cache:
    model_cache["proj_z"] = proj_z
```

--------------------------------

### Forward Mode Automatic Differentiation (JVP) with Segmented Polynomial

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/poly.html

This example demonstrates how to compute the Jacobian-Vector Product (JVP) for a Segmented Polynomial, which calculates the directional derivative. The resulting SP computes both the original values and their tangents.

```python
# Compute JVP with respect to input 0
sp_jvp, mapping = sp.jvp([True])
print(sp_jvp)
```

--------------------------------

### O3.__init__

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.O3.html

Initializes the O3 representation.

```APIDOC
## O3.__init__

### Description
Initializes the O3 representation.

### Method
```
__init__(_l_, _p_)
```

### Parameters
- **l** (_int_)
- **p** (_int_)

### Return Type
None
```

--------------------------------

### property dim

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.O3.html

Gets the dimension of the O3 representation.

```APIDOC
## property cuequivariance.O3.dim

### Description
Dimension of the representation.

### Returns
The dimension of the representation.

### Return type
`int`
```

--------------------------------

### RepArray Initialization and Usage

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_jax.RepArray.html

Demonstrates how to create and use RepArray instances with different initialization methods, including specifying representations directly or using default assumptions.

```APIDOC
## RepArray Initialization

### Description
Creates a `RepArray` by associating a `jax.Array` with group representations for its axes.

### Parameters
* **_reps_** (_dict_[_int_,_Rep_]|_Rep_|_Irreps_|_str_|_dict_[_int_,_Irreps_]|_dict_[_int_,_str_]) - The group representations for the axes.
* **_array_** (_Array_) - The underlying `jax.Array`.
* **_layout_** (_IrrepsLayout_|_None_) - Optional layout information for the representations.

### Example
```python
>>> import jax.numpy as jnp
>>> import cue
>>> import cuequivariance_jax as cuex

>>> # Example 1: Specifying Rep for each axis
>>> cuex.RepArray({0: cue.SO3(1), 1: cue.SO3(1)}, jnp.eye(3))
{0: 1, 1: 1}
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]

>>> # Example 2: Automatic conversion with assumed group and layout
>>> with cue.assume(cue.SO3, cue.ir_mul):
...     x = cuex.RepArray({0: "1", 1: "2"}, jnp.ones((3, 5)))
>>> x
{0: 1, 1: 2}
[[1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]]
>>> x.rep(0).irreps, x.rep(0).layout
(1, (irrep,mul))

>>> # Example 3: IrrepsArray initialization
>>> x = cuex.RepArray(
...     cue.Irreps("SO3", "2x0"), jnp.zeros((3, 2)), cue.ir_mul
... )
>>> x
{1: 2x0}
[[0. 0.]
 [0. 0.]
 [0. 0.]]

>>> # Example 4: Using default group and layout
>>> with cue.assume(cue.SO3, cue.ir_mul):
...     cuex.RepArray("2x0", jnp.array([1.0, 2.0]))
{0: 2x0} [1. 2.]
```
```

--------------------------------

### SegmentedOperand.num_segments

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedOperand.html

Gets the total number of segments currently in the operand.

```APIDOC
## SegmentedOperand.num_segments

### Description
Gets the number of segments in the operand.

### Returns
- int: The total number of segments.
```

--------------------------------

### SegmentedPolynomial Properties

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedPolynomial.html

Properties to get information about the number of operands for a SegmentedPolynomial.

```APIDOC
## SegmentedPolynomial Properties

### num_outputs

Get the number of output operands.

**Return type:**
`int`

### num_operands

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

**Return type:**
`int`
```

--------------------------------

### FullyConnectedTensorProduct Initialization

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.FullyConnectedTensorProduct.html

Initializes the FullyConnectedTensorProduct module with specified input and output irreps, and optional layout and weight configurations.

```APIDOC
## FullyConnectedTensorProduct

### Description
Initializes internal Module state, shared by both nn.Module and ScriptModule.

### Parameters
- **irreps_in1** (Irreps) - 
- **irreps_in2** (Irreps) - 
- **irreps_out** (Irreps) - 
- **layout** (IrrepsLayout | None) - 
- **layout_in1** (IrrepsLayout | None) - 
- **layout_in2** (IrrepsLayout | None) - 
- **layout_out** (IrrepsLayout | None) - 
- **shared_weights** (bool) - 
- **internal_weights** (bool) - 
- **device** (device | None) - 
- **dtype** (dtype | None) - 
- **math_dtype** (str | dtype | None) - 
- **use_fallback** (bool | None) - 
- **method** (str | None) -
```

--------------------------------

### get_segment_dimension

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedTensorProduct.html

Gets the dimension of a specific mode within a specific path of a segment.

```APIDOC
## get_segment_dimension(path, m)

### Description
Get the dimension of a specific mode in a specific path in a specific segment.

### Parameters
- **path** (int | Path) - The path identifier.
- **m** (str) - The name or identifier of the mode.

### Returns
- **int | set[int]** - The dimension of the specified mode in the specified path.
```

--------------------------------

### Get Dimension of SO3 Representation

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SO3.html

Retrieve the dimension of a given SO3 representation.

```python
SO3(1).dim
3
```

--------------------------------

### cuequivariance_torch.FullyConnectedTensorProduct

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the FullyConnectedTensorProduct class.

```APIDOC
## cuequivariance_torch.FullyConnectedTensorProduct.__init__()

### Description
Initializes the FullyConnectedTensorProduct class.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Get List of Multiplicities

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.Irreps.html

Access the list containing the multiplicities of each irreducible representation in the order they appear.

```python
>>> Irreps("SO3", "100x0 + 10x1").muls
[100, 10]
```

--------------------------------

### compressed_path_segment

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedTensorProduct.html

Calculates the starting indices of paths for the segments of a specified operand, assuming paths are sorted by that operand.

```APIDOC
## compressed_path_segment(operand)

### Description
Returns the starting indices of paths for the segments of the specified operand. This method is useful when paths are sorted by the specified operand.

### Parameters
- **operand** (int) - The index of the operand for which to find path groups.

### Returns
- **np.ndarray** - An array of starting path indices for each segment in the specified operand.

### Examples
```python
# Assuming 'indices' and 'operands' are defined elsewhere and structured appropriately
# indices[:, operand], operands[operand].num_segments -> compressed_path_segment(operand)
# Example outputs based on hypothetical inputs:
# [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]
```
```

--------------------------------

### Creating and Listing Operations

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.Operation.html

Demonstrates how to create a list of Operation objects, each defining a mapping of buffers to tensor product operands. The `list_to_string` method formats this list for display, showing the buffer mappings for each operation.

```python
>>> ops = \
...     Operation((0, 1, 2)),\
...     Operation((0, 1, 1, 2)),\
...     Operation((0, 1, 1, 1, 2)),
>>> print(Operation.list_to_string(ops, 2, 1))
(a, b) -> (C)
  a b C
  a b b C
  a b b b C
```

--------------------------------

### Iterate Over O3 Representations

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.O3.html

Obtain an iterator for all irreducible representations of the O3 group, starting with the trivial representation.

```python
>>> O3.iterator()
```

--------------------------------

### Get Dimensions Dictionary

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedTensorProduct.html

The `get_dimensions_dict` method returns a dictionary containing the dimensions of all modes within the tensor product.

```python
d.get_dimensions_dict()
```

--------------------------------

### assume.__init__

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.assume.html

Initializes the `assume` context manager or decorator with optional irrep class and layout.

```APIDOC
## assume.__init__

### Description
Initializes the `assume` context manager or decorator.

### Parameters
- **irrep_class** (Irrep | Type[Irrep] | str | None): The irrep class to assume. Can be an `Irrep` object, a type, a string representation, or None.
- **layout** (IrrepsLayout | None): The layout to assume. Can be an `IrrepsLayout` object or None.
```

--------------------------------

### Convert SegmentedTensorProduct to Base64

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedTensorProduct.html

Use `to_base64` to get a base64 encoded string representation of the descriptor. This is useful for serialization and transmission.

```python
print(d.to_base64())
```

--------------------------------

### cuequivariance.SO3

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the SO3 class.

```APIDOC
## cuequivariance.SO3.__init__()

### Description
Initializes the SO3 class.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Get Number of Irreducible Representations

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.Irreps.html

Obtain the total count of irreducible representations present in the Irreps object, summing up all multiplicities.

```python
>>> Irreps("SO3", "100x0 + 10x1").num_irreps
110
```

--------------------------------

### cuequivariance_torch.Inversion

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the Inversion class.

```APIDOC
## cuequivariance_torch.Inversion.__init__()

### Description
Initializes the Inversion class.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### cuequivariance.assume

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the assume module.

```APIDOC
## cuequivariance.assume.__init__()

### Description
Initializes the assume module.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Get Total Dimension of Representation

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.Irreps.html

Retrieve the total dimension of the representation, which is the sum of the dimensions of all irreducible components weighted by their multiplicities.

```python
>>> Irreps("SO3", "100x0 + 10x1").dim
130
```

--------------------------------

### Get Discrete Generators of SO3 Representation

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SO3.html

Compute the discrete generators of an SO3 representation, defined by \(\rho(n) = H^n\).

```python
rep.discrete_generators()
```

--------------------------------

### assume as Context Manager

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.assume.html

Demonstrates how to use `cue.assume` as a context manager to set the irrep class and layout within a `with` block.

```APIDOC
## assume as Context Manager

### Description
Use `cue.assume` as a context manager to temporarily set the irrep class and layout for a specific block of code.

### Usage
```python
with cue.assume(cue.SO3, cue.mul_ir):
    rep = cue.IrrepsAndLayout("2x1")
```

### Parameters
- **irrep_class**: The irrep class to assume (e.g., `cue.SO3`).
- **layout**: The layout to assume (e.g., `cue.mul_ir`).
```

--------------------------------

### cuequivariance.O3

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the O3 class.

```APIDOC
## cuequivariance.O3.__init__()

### Description
Initializes the O3 class.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Get O3 Regular Expression Pattern

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.O3.html

Retrieve the regular expression pattern used for parsing string representations of O3 irreps.

```python
>>> O3.regexp_pattern()
```

--------------------------------

### FullyConnectedTensorProduct

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.FullyConnectedTensorProduct.html

Initializes the FullyConnectedTensorProduct layer.

```APIDOC
## FullyConnectedTensorProduct

### Description
Fully connected tensor product layer.

### Parameters
- **irreps_in1** (Irreps) - Input irreps for the first operand.
- **irreps_in2** (Irreps) - Input irreps for the second operand.
- **irreps_out** (Irreps) - Output irreps.
- **layout** (IrrepsLayout, optional) - The layout of the input and output irreps. Default is `cue.mul_ir`.
- **layout_in1** (IrrepsLayout, optional) - The layout of the first input irreducible representations, by default `layout`.
- **layout_in2** (IrrepsLayout, optional) - The layout of the second input irreducible representations, by default `layout`.
- **layout_out** (IrrepsLayout, optional) - The layout of the output irreducible representations, by default `layout`.
- **shared_weights** (bool, optional) - Whether to share weights across the batch dimension. Default is True.
- **internal_weights** (bool, optional) - Whether to create module parameters for weights. Default is None.
- **device** (torch.device, optional) - The device to use for the operation.
- **dtype** (torch.dtype, optional) - The dtype to use for the operation weights.
- **math_dtype** (torch.dtype or string, optional) - The dtype to use for the math operations.
- **method** (str, optional) - The method to use for the linear layer, by default “fused_tp”.
- **use_fallback** (bool, optional, deprecated) - Whether to use a “fallback” implementation, now maps to method.
```

--------------------------------

### Get Algebra of SO3 Lie Group

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SO3.html

Retrieve the algebra of the SO3 Lie group, defined by the commutation relations \(\[X\_i, X\_j\].

```python
SO3.algebra()
```

--------------------------------

### Implement Indexed Linear with cuEquivariance (Naive Backend)

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html

Implement an indexed linear operation using cuEquivariance's Linear module with the 'naive' method. This method is more general but less performant than 'indexed_linear'. It requires transposing input features and manually rearranging weights.

```python
cue_irreps_in = cue.Irreps("O3", irreps_in)
cue_irreps_out = cue.Irreps("O3", irreps_out)

# The linear module
cue_lin = cuet.Linear(
    cue_irreps_in,
    cue_irreps_out,
    internal_weights=False,
    weight_classes=num_species,
    layout=cue.ir_mul,
    device=device,
    dtype=dtype,
    method='naive'
)

# Transposing the input
cue_in_feats = cuet.TransposeIrrepsLayout(
    cue_irreps_out,
    source=cue.mul_ir,
    target=cue.ir_mul,
    device=device,
    use_fallback=device=="cpu",
)(in_feats)

# Rearranging the weights by hand
cue_weight = skip_tp.weight.reshape(2*multiplicity, num_species, multiplicity
                ).transpose(0,1).reshape(num_species, -1)/np.sqrt(num_species)
# Performing the operation
cue_out_feats = cue_lin(cue_in_feats, weight=cue_weight, weight_indices=species)

print("Output shape:", cue_out_feats.shape)

# Transposing the output
cue_out_feats_transp = cuet.TransposeIrrepsLayout(
    cue_irreps_out,
    source=cue.ir_mul,
    target=cue.mul_ir,
    device=device,
    use_fallback=device=="cpu",
)(cue_out_feats)
# Comparing the result
print("Results match:", torch.allclose(out_feats, cue_out_feats_transp, atol=1e-3))
```

--------------------------------

### SymmetricContraction Operation with e3nn

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html

Performs a SymmetricContraction operation using e3nn. Requires one-hot encoded species. Used for initial setup and comparison.

```python
node_feats = torch.randn(num_nodes, multiplicity, irreps_in.dim // multiplicity, device=device, dtype=dtype)
species = torch.randint(0, num_species, (num_nodes,), device=device, dtype=torch.int64)
species_1hot = torch.nn.functional.one_hot(species, num_species).to(dtype).to(device)

# Perform operation
out_feats = sc(node_feats, species_1hot)

# Output shape
print("Output shape:", out_feats.shape)
```

--------------------------------

### FullyConnectedTensorProductConv Initialization

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.layers.FullyConnectedTensorProductConv.html

Initializes the FullyConnectedTensorProductConv layer with specified input, spherical harmonic, and output irreps, along with optional batch normalization, MLP configurations, layout, and fallback behavior.

```APIDOC
## FullyConnectedTensorProductConv(in_irreps, sh_irreps, out_irreps, batch_norm=True, mlp_channels=None, mlp_activation=nn.GELU(), layout=None, use_fallback=None)

### Description
Initializes the FullyConnectedTensorProductConv layer.

### Parameters
* **in_irreps** (_Irreps_) – Irreps for the input node features.
* **sh_irreps** (_Irreps_) – Irreps for the spherical harmonic representations of edge vectors.
* **out_irreps** (_Irreps_) – Irreps for the output.
* **batch_norm** (_bool_, optional) – If true, batch normalization is applied. Defaults to True.
* **mlp_channels** (Sequence of int, optional) – A sequence of integers defining the number of neurons in each layer in MLP before the output layer. If None, no MLP will be added. Defaults to None.
* **mlp_activation** (nn.Module or Sequence of nn.Module, optional) – A sequence of functions to be applied in between linear layers in MLP. Defaults to `nn.GELU()`.
* **layout** (IrrepsLayout, optional) – The layout of the input and output irreps. Default is `cue.mul_ir`.
* **use_fallback** (bool, optional) – If None (default), a CUDA kernel will be used if available. If False, a CUDA kernel will be used, and an exception is raised if it’s not available. If True, a PyTorch fallback method is used regardless of CUDA kernel availability.
```

--------------------------------

### Get Continuous Generators of O3 Representation

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.O3.html

Retrieve the continuous generators for a given O3 representation. These are used in the exponential map to obtain group elements.

```python
>>> O3(1, 1).continuous_generators()
```

--------------------------------

### Get Dimensions of a Specific Mode

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SegmentedTensorProduct.html

Use `get_dims` to retrieve the set of dimensions associated with a particular mode (e.g., 'u', 'v', 'w') in the tensor product.

```python
d.get_dims("u")
d.get_dims("v")
d.get_dims("w")
```

--------------------------------

### Performance Comparison: e3nn vs cuEquivariance

Source: https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html

Measures and prints the average execution time for both e3nn and cuEquivariance SymmetricContraction operations to demonstrate speedup. Includes synchronization for accurate CUDA timing.

```python
throwaway = 10 if device=="cuda" else 1
repetitions = 100 if device=="cuda" else 2

e3nn_times = []
for _ in range(throwaway):
    out_feats = sc(node_feats, species_1hot)
for _ in range(repetitions):
    if device=="cuda": torch.cuda.synchronize()
    t1 = time.perf_counter()
    out_feats = sc(node_feats, species_1hot)
    if device=="cuda": torch.cuda.synchronize()
    e3nn_times.append(time.perf_counter()-t1)

cuet_times = []
for _ in range(throwaway):
    cue_out_feats = cue_sc(cue_node_feats, species)
for _ in range(repetitions):
    if device=="cuda": torch.cuda.synchronize()
    t1 = time.perf_counter()
    cue_out_feats = cue_sc(cue_node_feats, species)
    if device=="cuda": torch.cuda.synchronize()
    cuet_times.append(time.perf_counter()-t1)
e3nn_avg = 1000*np.mean(e3nn_times)
cuet_avg = 1000*np.mean(cuet_times)
print(f"e3nn time: {e3nn_avg:.3} ms")
print(f"Cuequivariance time: {cuet_avg:.3} ms")
print(f"Speedup: {e3nn_avg/cuet_avg:.3}x")
```

--------------------------------

### Get Continuous Generators of SO3 Representation

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.SO3.html

Compute the continuous generators of an SO3 representation. These generators are used in the equation \(\rho(\alpha) = \exp\left(\alpha\_i X\_i\right)\).

```python
rep.continuous_generators()
```

--------------------------------

### cuequivariance.IrrepsAndLayout

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the IrrepsAndLayout class.

```APIDOC
## cuequivariance.IrrepsAndLayout.__init__()

### Description
Initializes the IrrepsAndLayout class.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### cuequivariance_torch.ChannelWiseTensorProduct

Source: https://docs.nvidia.com/cuda/cuequivariance/genindex.html

Initializes the ChannelWiseTensorProduct class.

```APIDOC
## cuequivariance_torch.ChannelWiseTensorProduct.__init__()

### Description
Initializes the ChannelWiseTensorProduct class.

### Method
__init__

### Endpoint
N/A (Python method)

### Parameters
None explicitly documented in the provided text.

### Request Example
N/A

### Response
N/A
```

--------------------------------

### Compute Spherical Harmonics with Segmented Polynomial

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_jax.segmented_polynomial.html

Example of computing spherical harmonics using the segmented_polynomial function. Requires importing cue and cuex, and defining the polynomial descriptor.

```python
>>> p = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2]).polynomial
>>> cuex.segmented_polynomial(

```

--------------------------------

### Initialize SymmetricContraction

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.SymmetricContraction.html

Instantiate the SymmetricContraction layer with specified input and output irreps, contraction degree, number of elements, layout, dtype, and device. Ensure CUDA is available for device selection.

```python
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
irreps_in = cue.Irreps("O3", "32x0e + 32x1o")
irreps_out = cue.Irreps("O3", "32x0e")
layer = SymmetricContraction(irreps_in, irreps_out, contraction_degree=3, num_elements=5, layout=cue.ir_mul, dtype=torch.float32, device=device)
```

--------------------------------

### Construct RepArray from Segments

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_jax.from_segments.html

Use cuex.from_segments to create a RepArray from specified irreps, segments, and shape. The example demonstrates constructing a RepArray with a specific SO3 irreps and segment data.

```python
cuex.from_segments(
    cue.Irreps("SO3", "2x0 + 1"),
    [jnp.array([[1.0], [2.0]]), jnp.array([[0.0], [0.0], [0.0]])],
    (-1,), cue.ir_mul)
{0: 2x0+1} [1. 2. 0. 0. 0.]
```

--------------------------------

### Import cuEquivariance Packages

Source: https://docs.nvidia.com/cuda/cuequivariance/index.html

Import the necessary cuEquivariance packages for JAX, PyTorch, or core non-ML components. Choose the appropriate import based on your framework.

```python
import cuequivariance as cue
# All the non-ML components

import cuequivariance_jax as cuex
# For the JAX implementations

import cuequivariance_torch as cuet
```

--------------------------------

### Get Slices for Representation Segments

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.Irreps.html

Generate a list of slice objects, where each slice corresponds to the range of indices occupied by a specific irreducible representation segment within the total representation.

```python
>>> Irreps("SO3", "100x0 + 10x1").slices()
[slice(0, 100, None), slice(100, 130, None)]
```

--------------------------------

### Creating IrrepsAndLayout with different layouts

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance.IrrepsAndLayout.html

Demonstrates how to create an IrrepsAndLayout object for a 2x1 representation of SO(3) using both 'ir_mul' and 'mul_ir' layouts. This is useful for understanding how data is organized for different group representations.

```python
import numpy as np
import cue

angles = np.array([np.pi, 0, 0])

with cue.assume("SO3", cue.ir_mul):
    rep = cue.IrrepsAndLayout("2x1")
R_ir_mul = rep.exp_map(angles, np.array([]))
```

```python
with cue.assume("SO3", cue.mul_ir):
    rep = cue.IrrepsAndLayout("2x1")
R_mul_ir = rep.exp_map(angles, np.array([]))
```

--------------------------------

### Using Repeats with jnp.repeat

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_jax.Repeats.html

Demonstrates how to instantiate the Repeats class and use it with jnp.repeat to control element repetition. Ensure JAX is imported as jnp.

```python
>>> a = Repeats(jnp.array([1, 0, 2]), 3)
>>> jnp.repeat(
...     jnp.array([0.1, 0.2, 0.3], dtype=jnp.float32),
...     a.repeats,
...     total_repeat_length=a.total_repeat_length,
... )
Array([0.1, 0.3, 0.3], dtype=float32)
```

--------------------------------

### FullyConnectedTensorProductConv Initialization

Source: https://docs.nvidia.com/cuda/cuequivariance/api/generated/cuequivariance_torch.layers.FullyConnectedTensorProductConv.html

Initializes the FullyConnectedTensorProductConv layer with specified input, spherical, and output representations, along with optional batch normalization and MLP configurations.

```APIDOC
## __init__

### Description
Initializes internal Module state.

### Parameters
* **in_irreps** (Irreps)
* **sh_irreps** (Irreps)
* **out_irreps** (Irreps)
* **batch_norm** (bool)
* **mlp_channels** (Sequence[int] | None)
* **mlp_activation** (Module | Sequence[Module] | None)
* **layout** (IrrepsLayout)
* **use_fallback** (bool | None)
```
Tokens:
Raw