cuequivariance-torch

e = cue.descriptors.channelwise_tensor_product(
    16 * cue.Irreps("SO3", "0 + 1"),
    cue.Irreps("SO3", "0 + 1"),
    cue.Irreps("SO3", "0 + 1"),
)
poly = e.polynomial
 
sp = cuet.SegmentedPolynomial(poly, method="uniform_1d")
 
w = torch.randn(1, poly.inputs[0].size, device="cuda")            # shared weights
x1 = torch.randn(batch, poly.inputs[1].size, device="cuda")       # batched input 1
x2 = torch.randn(batch, poly.inputs[2].size, device="cuda")       # batched input 2
 
# Basic forward
[out] = sp([w, x1, x2])
 
# With input gathering (e.g., gather x1 by node index)
senders = torch.randint(0, num_nodes, (num_edges,), device="cuda")
[out] = sp([w, x1, x2], input_indices={1: senders})
 
# With output scattering (accumulate into target nodes)
receivers = torch.randint(0, num_nodes, (num_edges,), device="cuda")
[out] = sp(
    [w, x1, x2],
    input_indices={1: senders},
    output_indices={0: receivers},
    output_shapes={0: torch.empty(num_nodes, 1, device="cuda")},
)

# Compute in float32 regardless of input dtype
sp = cuet.SegmentedPolynomial(poly, method="fused_tp", math_dtype=torch.float32)
 
# For fused_tp, math_dtype must be float32 or float64
# For naive, any torch.dtype works
# For uniform_1d, float32 or float64 (auto-selects float32 if input is e.g. float16)

# With internal weights (default: shared_weights=True, internal_weights=True)
tp = cuet.ChannelWiseTensorProduct(
    cue.Irreps("SO3", "32x0 + 32x1"),   # irreps_in1
    cue.Irreps("SO3", "0 + 1"),          # irreps_in2
    layout=cue.mul_ir,
    device="cuda",
    dtype=torch.float32,
)
# tp.weight_numel -- number of weight parameters
# tp.irreps_out -- output irreps (auto-computed)
 
x1 = torch.randn(batch, tp.irreps_in1.dim, device="cuda")
x2 = torch.randn(batch, tp.irreps_in2.dim, device="cuda")
 
out = tp(x1, x2)  # uses internal weight parameter
# out.shape == (batch, tp.irreps_out.dim)
 
# With external weights (shared_weights=False)
tp = cuet.ChannelWiseTensorProduct(
    cue.Irreps("SO3", "32x0 + 32x1"),
    cue.Irreps("SO3", "0 + 1"),
    layout=cue.mul_ir,
    shared_weights=False,
    device="cuda",
)
w = torch.randn(batch, tp.weight_numel, device="cuda")
out = tp(x1, x2, weight=w)
 
# With gather/scatter for graph neural networks
out = tp(x1, x2, weight=w, indices_1=senders, indices_out=receivers, size_out=num_nodes)

tp = cuet.FullyConnectedTensorProduct(
    cue.Irreps("O3", "4x0e + 4x1o"),    # irreps_in1
    cue.Irreps("O3", "0e + 1o"),         # irreps_in2
    cue.Irreps("O3", "4x0e + 4x1o"),    # irreps_out
    layout=cue.mul_ir,
    internal_weights=True,
    device="cuda",
)
 
out = tp(x1, x2)  # uses internal weights
# or: out = tp(x1, x2, weight=w)  # external weights

linear = cuet.Linear(
    cue.Irreps("SO3", "4x0 + 2x1"),     # irreps_in
    cue.Irreps("SO3", "3x0 + 5x1"),     # irreps_out
    layout=cue.mul_ir,
    internal_weights=True,
    device="cuda",
)
 
out = linear(x)
 
# Species-indexed weights (different weights per atom type)
linear = cuet.Linear(
    irreps_in, irreps_out,
    weight_classes=50,   # 50 different weight sets
    internal_weights=True,
    device="cuda",
)
out = linear(x, weight_indices=species_indices)  # species_indices: (batch,) int tensor

sc = cuet.SymmetricContraction(
    cue.Irreps("O3", "32x0e + 32x1o"),  # irreps_in (uniform mul required)
    cue.Irreps("O3", "32x0e"),           # irreps_out (uniform mul required)
    contraction_degree=3,
    num_elements=95,                      # number of chemical elements
    layout=cue.ir_mul,
    dtype=torch.float32,
    device="cuda",
)
 
# indices: (batch,) int tensor selecting which element weights to use
out = sc(x, indices)
# out.shape == (batch, irreps_out.dim)

sh = cuet.SphericalHarmonics(
    ls=[0, 1, 2, 3],
    normalize=True,
    device="cuda",
)
 
vectors = torch.randn(batch, 3, device="cuda")
out = sh(vectors)
# out.shape == (batch, 1 + 3 + 5 + 7)  -- sum of 2l+1

# Rotation (SO3 or O3 irreps)
rot = cuet.Rotation(
    cue.Irreps("SO3", "4x0 + 2x1 + 1x2"),
    layout=cue.ir_mul,
    device="cuda",
)
 
# Euler angles (YXY convention)
gamma = torch.tensor([0.1], device="cuda")
beta = torch.tensor([0.2], device="cuda")
alpha = torch.tensor([0.3], device="cuda")
out = rot(gamma, beta, alpha, x)
 
# Helper: encode angle for rotation
encoded = cuet.encode_rotation_angle(angle, ell=3)  # cos/sin encoding
 
# Helper: 3D vector to Euler angles
beta, alpha = cuet.vector_to_euler_angles(vector)
 
# Inversion (O3 irreps only)
inv = cuet.Inversion(
    cue.Irreps("O3", "4x0e + 2x1o"),
    layout=cue.ir_mul,
    device="cuda",
)
out = inv(x)

bn = cuet.layers.BatchNorm(
    cue.Irreps("O3", "4x0e + 4x1o"),
    layout=cue.mul_ir,
    eps=1e-5,
    momentum=0.1,
    affine=True,
)
 
out = bn(x)  # x.shape == (batch, irreps.dim)

conv = cuet.layers.FullyConnectedTensorProductConv(
    in_irreps=cue.Irreps("O3", "4x0e + 4x1o"),
    sh_irreps=cue.Irreps("O3", "0e + 1o"),
    out_irreps=cue.Irreps("O3", "4x0e + 4x1o"),
    mlp_channels=[16, 32, 32],
    mlp_activation=torch.nn.ReLU(),
    batch_norm=True,
    layout=cue.ir_mul,
)
 
# graph = ((src, dst), (num_src_nodes, num_dst_nodes))
graph = ((src, dst), (num_src_nodes, num_dst_nodes))
 
out = conv(src_features, edge_sh, edge_emb, graph, reduce="mean")
# out.shape == (num_dst_nodes, out_irreps.dim)
 
# Optional: separate scalar features for efficient first-layer GEMM
out = conv(src_features, edge_sh, edge_emb, graph,
           src_scalars=src_scalars, dst_scalars=dst_scalars)

# Triangle attention with pair bias
kv_lengths = cuet.mask_to_kv_lengths(prefix_mask)
out = cuet.triangle_attention(q, k, v, bias, scale=scale, kv_lengths=kv_lengths)
# Arbitrary dense masks are supported through the fallback path:
out = cuet.triangle_attention(q, k, v, bias, mask=holey_mask, scale=scale)
# q, k, v: (B, N, H, Q/K, D), bias: (B, 1, H, Q, K)
 
# Triangle multiplicative update
out = cuet.triangle_multiplicative_update(
    x,         # (B, I, J, C)
    mask=mask, # (B, I, J)
    precision=cuet.TriMulPrecision.DEFAULT,
)
 
# Attention with pair bias (strict OpenFold3/Boltz contract)
out, _ = cuet.attention_pair_bias(
    single_repr, pair_repr, mask, num_heads,
    w_ln_a, b_ln_a,
    w_proj_q, b_proj_q, w_proj_k, w_proj_v,
    w_proj_g, w_proj_o, w_proj_z,
    w_ln_z=w_ln_z, b_ln_z=b_ln_z,
)
# Optional generalized projection params are keyword-only:
# b_proj_k, b_proj_v, w_ln_q, b_ln_q, w_ln_k, b_ln_k, b_proj_g (default None),
# and norm_kind (default "layer_norm").

# ONNX export
table = cuet.onnx_custom_translation_table()
onnx_program = torch.onnx.export(model, inputs, custom_translation_table=table)
 
# TensorRT plugin registration
cuet.register_tensorrt_plugins()

import torch
import cuequivariance as cue
import cuequivariance_torch as cuet
 
class SimpleGNN(torch.nn.Module):
    def __init__(self, irreps_in, irreps_sh, irreps_out):
        super().__init__()
        self.tp = cuet.ChannelWiseTensorProduct(
            irreps_in, irreps_sh, layout=cue.mul_ir,
            shared_weights=False, device="cuda",
        )
        self.linear = cuet.Linear(
            self.tp.irreps_out, irreps_out,
            layout=cue.mul_ir, internal_weights=True, device="cuda",
        )
        self.sh = cuet.SphericalHarmonics(
            ls=[ir.l for _, ir in irreps_sh], normalize=True, device="cuda",
        )
 
    def forward(self, node_feats, edge_vec, edge_index, num_nodes):
        src, dst = edge_index
        edge_sh = self.sh(edge_vec)
        w = torch.randn(1, self.tp.weight_numel, device=node_feats.device)
 
        # Message: tensor product on edges with scatter to destination nodes
        messages = self.tp(
            node_feats, edge_sh, weight=w,
            indices_1=src, indices_2=None,
            indices_out=dst, size_out=num_nodes,
        )
        return self.linear(messages)