[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"skill-nvidia-cuequivariance-torch":3,"mdc-q0kug9-key":37,"related-repo-nvidia-cuequivariance-torch":3101,"related-org-nvidia-cuequivariance-torch":3135},{"slug":4,"name":4,"fn":5,"description":6,"org":7,"tags":11,"stars":26,"repoUrl":27,"updatedAt":28,"license":29,"forks":30,"topics":31,"repo":32,"sourceUrl":35,"mdContent":36},"cuequivariance-torch","execute equivariant tensor products in PyTorch","Execute equivariant tensor products in PyTorch using SegmentedPolynomial (naive\u002Funiform_1d\u002Ffused_tp\u002Findexed_linear), high-level operations (ChannelWiseTensorProduct, FullyConnectedTensorProduct, Linear, SymmetricContraction, SphericalHarmonics, Rotation), and layers (BatchNorm, FullyConnectedTensorProductConv). Use when writing PyTorch code with cuequivariance.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},"nvidia","NVIDIA","https:\u002F\u002Fpexgzepcugksgbtrxkhf.supabase.co\u002Fstorage\u002Fv1\u002Fobject\u002Fpublic\u002Forg-logos\u002Fnvidia.png",[12,16,19,22,23],{"name":13,"slug":14,"type":15},"Deep Learning","deep-learning","tag",{"name":17,"slug":18,"type":15},"Mathematics","mathematics",{"name":20,"slug":21,"type":15},"PyTorch","pytorch",{"name":9,"slug":8,"type":15},{"name":24,"slug":25,"type":15},"Physics","physics",409,"https:\u002F\u002Fgithub.com\u002FNVIDIA\u002FcuEquivariance","2026-07-23T05:43:52.496117",null,36,[],{"repoUrl":27,"stars":26,"forks":30,"topics":33,"description":34},[],"cuEquivariance is a math library that is a collective of low-level primitives and tensor ops to accelerate widely-used models, like DiffDock, MACE, Allegro and NEQUIP, based on equivariant neural networks. Also includes kernels for accelerated structure prediction.","https:\u002F\u002Fgithub.com\u002FNVIDIA\u002FcuEquivariance\u002Ftree\u002FHEAD\u002Fcuequivariance_torch\u002Fcuequivariance_torch","---\nname: cuequivariance-torch\ndescription: Execute equivariant tensor products in PyTorch using SegmentedPolynomial (naive\u002Funiform_1d\u002Ffused_tp\u002Findexed_linear), high-level operations (ChannelWiseTensorProduct, FullyConnectedTensorProduct, Linear, SymmetricContraction, SphericalHarmonics, Rotation), and layers (BatchNorm, FullyConnectedTensorProductConv). Use when writing PyTorch code with cuequivariance.\n---\n\n# cuequivariance_torch: Executing Equivariant Polynomials in PyTorch\n\n## Overview\n\n`cuequivariance_torch` (imported as `cuet`) executes `cuequivariance` polynomials on GPU via PyTorch. It provides:\n\n1. **Core primitive**: `cuet.SegmentedPolynomial` — `torch.nn.Module` with multiple CUDA backends\n2. **High-level operations** (`torch.nn.Module`): `ChannelWiseTensorProduct`, `FullyConnectedTensorProduct`, `Linear`, `SymmetricContraction`, `SphericalHarmonics`, `Rotation`, `Inversion`\n3. **Layers**: `cuet.layers.BatchNorm`, `cuet.layers.FullyConnectedTensorProductConv` (message passing)\n4. **Utilities**: `triangle_attention`, `triangle_multiplicative_update`, `attention_pair_bias` (AlphaFold3-style)\n5. **Export support**: `onnx_custom_translation_table()`, `register_tensorrt_plugins()`\n\n## Execution methods\n\n| Method | Backend | Requirements |\n|--------|---------|-------------|\n| `\"naive\"` | Pure PyTorch (einsum) | Always works, any platform |\n| `\"uniform_1d\"` | CUDA kernel | GPU, all segments uniform shape within each operand, single mode |\n| `\"fused_tp\"` | CUDA kernel | GPU, 3- or 4-operand contractions, float32\u002Ffloat64 |\n| `\"indexed_linear\"` | CUDA kernel | GPU, linear with indexed weights, sorted indices |\n\n## Core primitive: SegmentedPolynomial\n\n```python\nimport torch\nimport cuequivariance as cue\nimport cuequivariance_torch as cuet\n\n# Build a descriptor\ne = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2])\npoly = e.polynomial\n\n# Create the module\nsp = cuet.SegmentedPolynomial(poly, method=\"uniform_1d\")\n\n# Forward pass\nx = torch.randn(batch, 3, device=\"cuda\")\n[output] = sp([x])\n# output.shape == (batch, 9)  -- 1 + 3 + 5\n```\n\n### Inputs, indexing, and scatter\n\n```python\ne = cue.descriptors.channelwise_tensor_product(\n    16 * cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n)\npoly = e.polynomial\n\nsp = cuet.SegmentedPolynomial(poly, method=\"uniform_1d\")\n\nw = torch.randn(1, poly.inputs[0].size, device=\"cuda\")            # shared weights\nx1 = torch.randn(batch, poly.inputs[1].size, device=\"cuda\")       # batched input 1\nx2 = torch.randn(batch, poly.inputs[2].size, device=\"cuda\")       # batched input 2\n\n# Basic forward\n[out] = sp([w, x1, x2])\n\n# With input gathering (e.g., gather x1 by node index)\nsenders = torch.randint(0, num_nodes, (num_edges,), device=\"cuda\")\n[out] = sp([w, x1, x2], input_indices={1: senders})\n\n# With output scattering (accumulate into target nodes)\nreceivers = torch.randint(0, num_nodes, (num_edges,), device=\"cuda\")\n[out] = sp(\n    [w, x1, x2],\n    input_indices={1: senders},\n    output_indices={0: receivers},\n    output_shapes={0: torch.empty(num_nodes, 1, device=\"cuda\")},\n)\n```\n\n### Math dtype control\n\n```python\n# Compute in float32 regardless of input dtype\nsp = cuet.SegmentedPolynomial(poly, method=\"fused_tp\", math_dtype=torch.float32)\n\n# For fused_tp, math_dtype must be float32 or float64\n# For naive, any torch.dtype works\n# For uniform_1d, float32 or float64 (auto-selects float32 if input is e.g. float16)\n```\n\n## High-level operations\n\nAll operations are `torch.nn.Module` subclasses. They wrap `SegmentedPolynomial` and handle layout transposition automatically.\n\n### Memory layout\n\n`IrrepsLayout` controls memory order within each `(mul, ir)` block:\n\n- `cue.mul_ir`: data ordered as `(mul, ir.dim)` — **default, compatible with e3nn**\n- `cue.ir_mul`: data ordered as `(ir.dim, mul)` — **used internally by descriptors**\n\nOperations accept `layout` (applies to all), or per-operand `layout_in1`, `layout_in2`, `layout_out`.\n\n### ChannelWiseTensorProduct\n\nChannel-wise tensor product: pairs channels of `x1` with channels of `x2`.\n\n```python\n# With internal weights (default: shared_weights=True, internal_weights=True)\ntp = cuet.ChannelWiseTensorProduct(\n    cue.Irreps(\"SO3\", \"32x0 + 32x1\"),   # irreps_in1\n    cue.Irreps(\"SO3\", \"0 + 1\"),          # irreps_in2\n    layout=cue.mul_ir,\n    device=\"cuda\",\n    dtype=torch.float32,\n)\n# tp.weight_numel -- number of weight parameters\n# tp.irreps_out -- output irreps (auto-computed)\n\nx1 = torch.randn(batch, tp.irreps_in1.dim, device=\"cuda\")\nx2 = torch.randn(batch, tp.irreps_in2.dim, device=\"cuda\")\n\nout = tp(x1, x2)  # uses internal weight parameter\n# out.shape == (batch, tp.irreps_out.dim)\n\n# With external weights (shared_weights=False)\ntp = cuet.ChannelWiseTensorProduct(\n    cue.Irreps(\"SO3\", \"32x0 + 32x1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n    layout=cue.mul_ir,\n    shared_weights=False,\n    device=\"cuda\",\n)\nw = torch.randn(batch, tp.weight_numel, device=\"cuda\")\nout = tp(x1, x2, weight=w)\n\n# With gather\u002Fscatter for graph neural networks\nout = tp(x1, x2, weight=w, indices_1=senders, indices_out=receivers, size_out=num_nodes)\n```\n\nDefault method: `\"uniform_1d\"` if segments are uniform, else `\"naive\"`.\n\n### FullyConnectedTensorProduct\n\nAll input irrep pairs contribute to all output irreps (dense contraction).\n\n```python\ntp = cuet.FullyConnectedTensorProduct(\n    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),    # irreps_in1\n    cue.Irreps(\"O3\", \"0e + 1o\"),         # irreps_in2\n    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),    # irreps_out\n    layout=cue.mul_ir,\n    internal_weights=True,\n    device=\"cuda\",\n)\n\nout = tp(x1, x2)  # uses internal weights\n# or: out = tp(x1, x2, weight=w)  # external weights\n```\n\nDefault method: `\"fused_tp\"`.\n\n### Linear\n\nEquivariant linear layer (weight-only, no second input).\n\n```python\nlinear = cuet.Linear(\n    cue.Irreps(\"SO3\", \"4x0 + 2x1\"),     # irreps_in\n    cue.Irreps(\"SO3\", \"3x0 + 5x1\"),     # irreps_out\n    layout=cue.mul_ir,\n    internal_weights=True,\n    device=\"cuda\",\n)\n\nout = linear(x)\n\n# Species-indexed weights (different weights per atom type)\nlinear = cuet.Linear(\n    irreps_in, irreps_out,\n    weight_classes=50,   # 50 different weight sets\n    internal_weights=True,\n    device=\"cuda\",\n)\nout = linear(x, weight_indices=species_indices)  # species_indices: (batch,) int tensor\n```\n\nDefault method: `\"naive\"`. Use `method=\"fused_tp\"` for CUDA acceleration.\n\n### SymmetricContraction\n\nMACE-style symmetric contraction with element-indexed weights.\n\n```python\nsc = cuet.SymmetricContraction(\n    cue.Irreps(\"O3\", \"32x0e + 32x1o\"),  # irreps_in (uniform mul required)\n    cue.Irreps(\"O3\", \"32x0e\"),           # irreps_out (uniform mul required)\n    contraction_degree=3,\n    num_elements=95,                      # number of chemical elements\n    layout=cue.ir_mul,\n    dtype=torch.float32,\n    device=\"cuda\",\n)\n\n# indices: (batch,) int tensor selecting which element weights to use\nout = sc(x, indices)\n# out.shape == (batch, irreps_out.dim)\n```\n\nDefault method: `\"uniform_1d\"` if segments are uniform, else `\"naive\"`.\n\n### SphericalHarmonics\n\n```python\nsh = cuet.SphericalHarmonics(\n    ls=[0, 1, 2, 3],\n    normalize=True,\n    device=\"cuda\",\n)\n\nvectors = torch.randn(batch, 3, device=\"cuda\")\nout = sh(vectors)\n# out.shape == (batch, 1 + 3 + 5 + 7)  -- sum of 2l+1\n```\n\nDefault method: `\"uniform_1d\"`.\n\n### Rotation and Inversion\n\n```python\n# Rotation (SO3 or O3 irreps)\nrot = cuet.Rotation(\n    cue.Irreps(\"SO3\", \"4x0 + 2x1 + 1x2\"),\n    layout=cue.ir_mul,\n    device=\"cuda\",\n)\n\n# Euler angles (YXY convention)\ngamma = torch.tensor([0.1], device=\"cuda\")\nbeta = torch.tensor([0.2], device=\"cuda\")\nalpha = torch.tensor([0.3], device=\"cuda\")\nout = rot(gamma, beta, alpha, x)\n\n# Helper: encode angle for rotation\nencoded = cuet.encode_rotation_angle(angle, ell=3)  # cos\u002Fsin encoding\n\n# Helper: 3D vector to Euler angles\nbeta, alpha = cuet.vector_to_euler_angles(vector)\n\n# Inversion (O3 irreps only)\ninv = cuet.Inversion(\n    cue.Irreps(\"O3\", \"4x0e + 2x1o\"),\n    layout=cue.ir_mul,\n    device=\"cuda\",\n)\nout = inv(x)\n```\n\n## Layers\n\n### BatchNorm\n\nBatch normalization for equivariant representations (adapted from e3nn).\n\n```python\nbn = cuet.layers.BatchNorm(\n    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n    layout=cue.mul_ir,\n    eps=1e-5,\n    momentum=0.1,\n    affine=True,\n)\n\nout = bn(x)  # x.shape == (batch, irreps.dim)\n```\n\n### FullyConnectedTensorProductConv\n\nMessage passing layer for equivariant GNNs (DiffDock-style).\n\n```python\nconv = cuet.layers.FullyConnectedTensorProductConv(\n    in_irreps=cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n    sh_irreps=cue.Irreps(\"O3\", \"0e + 1o\"),\n    out_irreps=cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n    mlp_channels=[16, 32, 32],\n    mlp_activation=torch.nn.ReLU(),\n    batch_norm=True,\n    layout=cue.ir_mul,\n)\n\n# graph = ((src, dst), (num_src_nodes, num_dst_nodes))\ngraph = ((src, dst), (num_src_nodes, num_dst_nodes))\n\nout = conv(src_features, edge_sh, edge_emb, graph, reduce=\"mean\")\n# out.shape == (num_dst_nodes, out_irreps.dim)\n\n# Optional: separate scalar features for efficient first-layer GEMM\nout = conv(src_features, edge_sh, edge_emb, graph,\n           src_scalars=src_scalars, dst_scalars=dst_scalars)\n```\n\n## Triangle operations (AlphaFold2-style)\n\nRequire `cuequivariance_ops_torch`.\n\n```python\n# Triangle attention with pair bias\nkv_lengths = cuet.mask_to_kv_lengths(prefix_mask)\nout = cuet.triangle_attention(q, k, v, bias, scale=scale, kv_lengths=kv_lengths)\n# Arbitrary dense masks are supported through the fallback path:\nout = cuet.triangle_attention(q, k, v, bias, mask=holey_mask, scale=scale)\n# q, k, v: (B, N, H, Q\u002FK, D), bias: (B, 1, H, Q, K)\n\n# Triangle multiplicative update\nout = cuet.triangle_multiplicative_update(\n    x,         # (B, I, J, C)\n    mask=mask, # (B, I, J)\n    precision=cuet.TriMulPrecision.DEFAULT,\n)\n\n# Attention with pair bias (strict OpenFold3\u002FBoltz contract)\nout, _ = cuet.attention_pair_bias(\n    single_repr, pair_repr, mask, num_heads,\n    w_ln_a, b_ln_a,\n    w_proj_q, b_proj_q, w_proj_k, w_proj_v,\n    w_proj_g, w_proj_o, w_proj_z,\n    w_ln_z=w_ln_z, b_ln_z=b_ln_z,\n)\n# Optional generalized (Proteina\u002FComplexa) params, all keyword-only and None by\n# default: b_proj_k, b_proj_v, w_ln_q, b_ln_q, w_ln_k, b_ln_k, b_proj_g.\n```\n\n## ONNX and TensorRT export\n\n```python\n# ONNX export\ntable = cuet.onnx_custom_translation_table()\nonnx_program = torch.onnx.export(model, inputs, custom_translation_table=table)\n\n# TensorRT plugin registration\ncuet.register_tensorrt_plugins()\n```\n\n## Complete GNN example\n\n```python\nimport torch\nimport cuequivariance as cue\nimport cuequivariance_torch as cuet\n\nclass SimpleGNN(torch.nn.Module):\n    def __init__(self, irreps_in, irreps_sh, irreps_out):\n        super().__init__()\n        self.tp = cuet.ChannelWiseTensorProduct(\n            irreps_in, irreps_sh, layout=cue.mul_ir,\n            shared_weights=False, device=\"cuda\",\n        )\n        self.linear = cuet.Linear(\n            self.tp.irreps_out, irreps_out,\n            layout=cue.mul_ir, internal_weights=True, device=\"cuda\",\n        )\n        self.sh = cuet.SphericalHarmonics(\n            ls=[ir.l for _, ir in irreps_sh], normalize=True, device=\"cuda\",\n        )\n\n    def forward(self, node_feats, edge_vec, edge_index, num_nodes):\n        src, dst = edge_index\n        edge_sh = self.sh(edge_vec)\n        w = torch.randn(1, self.tp.weight_numel, device=node_feats.device)\n\n        # Message: tensor product on edges with scatter to destination nodes\n        messages = self.tp(\n            node_feats, edge_sh, weight=w,\n            indices_1=src, indices_2=None,\n            indices_out=dst, size_out=num_nodes,\n        )\n        return self.linear(messages)\n```\n\n## Key file locations\n\n| Component | Path |\n|-----------|------|\n| `SegmentedPolynomial` | `cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial.py` |\n| `uniform_1d` backend | `cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_uniform_1d.py` |\n| `naive` backend | `cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_naive.py` |\n| `fused_tp` backend | `cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_fused_tp.py` |\n| `indexed_linear` backend | `cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_indexed_linear.py` |\n| `ChannelWiseTensorProduct` | `cuequivariance_torch\u002Foperations\u002Ftp_channel_wise.py` |\n| `FullyConnectedTensorProduct` | `cuequivariance_torch\u002Foperations\u002Ftp_fully_connected.py` |\n| `Linear` | `cuequivariance_torch\u002Foperations\u002Flinear.py` |\n| `SymmetricContraction` | `cuequivariance_torch\u002Foperations\u002Fsymmetric_contraction.py` |\n| `SphericalHarmonics` | `cuequivariance_torch\u002Foperations\u002Fspherical_harmonics.py` |\n| `Rotation` \u002F `Inversion` | `cuequivariance_torch\u002Foperations\u002Frotation.py` |\n| `BatchNorm` | `cuequivariance_torch\u002Flayers\u002Fbatchnorm.py` |\n| `FullyConnectedTensorProductConv` | `cuequivariance_torch\u002Flayers\u002Ftp_conv_fully_connected.py` |\n| Triangle operations | `cuequivariance_torch\u002Fprimitives\u002Ftriangle.py` |\n| Layout transposition | `cuequivariance_torch\u002Fprimitives\u002Ftranspose.py` |\n",{"data":38,"body":39},{"name":4,"description":6},{"type":40,"children":41},"root",[42,51,58,87,262,268,387,393,540,547,782,788,842,847,867,873,892,941,976,981,1001,1240,1258,1263,1268,1359,1369,1374,1379,1520,1539,1544,1549,1656,1672,1677,1753,1763,1769,1974,1979,1985,1990,2066,2072,2077,2231,2237,2249,2444,2450,2504,2510,2757,2763,3095],{"type":43,"tag":44,"props":45,"children":47},"element","h1",{"id":46},"cuequivariance_torch-executing-equivariant-polynomials-in-pytorch",[48],{"type":49,"value":50},"text","cuequivariance_torch: Executing Equivariant Polynomials in PyTorch",{"type":43,"tag":52,"props":53,"children":55},"h2",{"id":54},"overview",[56],{"type":49,"value":57},"Overview",{"type":43,"tag":59,"props":60,"children":61},"p",{},[62,69,71,77,79,85],{"type":43,"tag":63,"props":64,"children":66},"code",{"className":65},[],[67],{"type":49,"value":68},"cuequivariance_torch",{"type":49,"value":70}," (imported as ",{"type":43,"tag":63,"props":72,"children":74},{"className":73},[],[75],{"type":49,"value":76},"cuet",{"type":49,"value":78},") executes ",{"type":43,"tag":63,"props":80,"children":82},{"className":81},[],[83],{"type":49,"value":84},"cuequivariance",{"type":49,"value":86}," polynomials on GPU via PyTorch. It provides:",{"type":43,"tag":88,"props":89,"children":90},"ol",{},[91,119,185,209,240],{"type":43,"tag":92,"props":93,"children":94},"li",{},[95,101,103,109,111,117],{"type":43,"tag":96,"props":97,"children":98},"strong",{},[99],{"type":49,"value":100},"Core primitive",{"type":49,"value":102},": ",{"type":43,"tag":63,"props":104,"children":106},{"className":105},[],[107],{"type":49,"value":108},"cuet.SegmentedPolynomial",{"type":49,"value":110}," — ",{"type":43,"tag":63,"props":112,"children":114},{"className":113},[],[115],{"type":49,"value":116},"torch.nn.Module",{"type":49,"value":118}," with multiple CUDA backends",{"type":43,"tag":92,"props":120,"children":121},{},[122,127,129,134,136,142,144,150,151,157,158,164,165,171,172,178,179],{"type":43,"tag":96,"props":123,"children":124},{},[125],{"type":49,"value":126},"High-level operations",{"type":49,"value":128}," (",{"type":43,"tag":63,"props":130,"children":132},{"className":131},[],[133],{"type":49,"value":116},{"type":49,"value":135},"): ",{"type":43,"tag":63,"props":137,"children":139},{"className":138},[],[140],{"type":49,"value":141},"ChannelWiseTensorProduct",{"type":49,"value":143},", ",{"type":43,"tag":63,"props":145,"children":147},{"className":146},[],[148],{"type":49,"value":149},"FullyConnectedTensorProduct",{"type":49,"value":143},{"type":43,"tag":63,"props":152,"children":154},{"className":153},[],[155],{"type":49,"value":156},"Linear",{"type":49,"value":143},{"type":43,"tag":63,"props":159,"children":161},{"className":160},[],[162],{"type":49,"value":163},"SymmetricContraction",{"type":49,"value":143},{"type":43,"tag":63,"props":166,"children":168},{"className":167},[],[169],{"type":49,"value":170},"SphericalHarmonics",{"type":49,"value":143},{"type":43,"tag":63,"props":173,"children":175},{"className":174},[],[176],{"type":49,"value":177},"Rotation",{"type":49,"value":143},{"type":43,"tag":63,"props":180,"children":182},{"className":181},[],[183],{"type":49,"value":184},"Inversion",{"type":43,"tag":92,"props":186,"children":187},{},[188,193,194,200,201,207],{"type":43,"tag":96,"props":189,"children":190},{},[191],{"type":49,"value":192},"Layers",{"type":49,"value":102},{"type":43,"tag":63,"props":195,"children":197},{"className":196},[],[198],{"type":49,"value":199},"cuet.layers.BatchNorm",{"type":49,"value":143},{"type":43,"tag":63,"props":202,"children":204},{"className":203},[],[205],{"type":49,"value":206},"cuet.layers.FullyConnectedTensorProductConv",{"type":49,"value":208}," (message passing)",{"type":43,"tag":92,"props":210,"children":211},{},[212,217,218,224,225,231,232,238],{"type":43,"tag":96,"props":213,"children":214},{},[215],{"type":49,"value":216},"Utilities",{"type":49,"value":102},{"type":43,"tag":63,"props":219,"children":221},{"className":220},[],[222],{"type":49,"value":223},"triangle_attention",{"type":49,"value":143},{"type":43,"tag":63,"props":226,"children":228},{"className":227},[],[229],{"type":49,"value":230},"triangle_multiplicative_update",{"type":49,"value":143},{"type":43,"tag":63,"props":233,"children":235},{"className":234},[],[236],{"type":49,"value":237},"attention_pair_bias",{"type":49,"value":239}," (AlphaFold3-style)",{"type":43,"tag":92,"props":241,"children":242},{},[243,248,249,255,256],{"type":43,"tag":96,"props":244,"children":245},{},[246],{"type":49,"value":247},"Export support",{"type":49,"value":102},{"type":43,"tag":63,"props":250,"children":252},{"className":251},[],[253],{"type":49,"value":254},"onnx_custom_translation_table()",{"type":49,"value":143},{"type":43,"tag":63,"props":257,"children":259},{"className":258},[],[260],{"type":49,"value":261},"register_tensorrt_plugins()",{"type":43,"tag":52,"props":263,"children":265},{"id":264},"execution-methods",[266],{"type":49,"value":267},"Execution methods",{"type":43,"tag":269,"props":270,"children":271},"table",{},[272,296],{"type":43,"tag":273,"props":274,"children":275},"thead",{},[276],{"type":43,"tag":277,"props":278,"children":279},"tr",{},[280,286,291],{"type":43,"tag":281,"props":282,"children":283},"th",{},[284],{"type":49,"value":285},"Method",{"type":43,"tag":281,"props":287,"children":288},{},[289],{"type":49,"value":290},"Backend",{"type":43,"tag":281,"props":292,"children":293},{},[294],{"type":49,"value":295},"Requirements",{"type":43,"tag":297,"props":298,"children":299},"tbody",{},[300,323,345,366],{"type":43,"tag":277,"props":301,"children":302},{},[303,313,318],{"type":43,"tag":304,"props":305,"children":306},"td",{},[307],{"type":43,"tag":63,"props":308,"children":310},{"className":309},[],[311],{"type":49,"value":312},"\"naive\"",{"type":43,"tag":304,"props":314,"children":315},{},[316],{"type":49,"value":317},"Pure PyTorch (einsum)",{"type":43,"tag":304,"props":319,"children":320},{},[321],{"type":49,"value":322},"Always works, any platform",{"type":43,"tag":277,"props":324,"children":325},{},[326,335,340],{"type":43,"tag":304,"props":327,"children":328},{},[329],{"type":43,"tag":63,"props":330,"children":332},{"className":331},[],[333],{"type":49,"value":334},"\"uniform_1d\"",{"type":43,"tag":304,"props":336,"children":337},{},[338],{"type":49,"value":339},"CUDA kernel",{"type":43,"tag":304,"props":341,"children":342},{},[343],{"type":49,"value":344},"GPU, all segments uniform shape within each operand, single mode",{"type":43,"tag":277,"props":346,"children":347},{},[348,357,361],{"type":43,"tag":304,"props":349,"children":350},{},[351],{"type":43,"tag":63,"props":352,"children":354},{"className":353},[],[355],{"type":49,"value":356},"\"fused_tp\"",{"type":43,"tag":304,"props":358,"children":359},{},[360],{"type":49,"value":339},{"type":43,"tag":304,"props":362,"children":363},{},[364],{"type":49,"value":365},"GPU, 3- or 4-operand contractions, float32\u002Ffloat64",{"type":43,"tag":277,"props":367,"children":368},{},[369,378,382],{"type":43,"tag":304,"props":370,"children":371},{},[372],{"type":43,"tag":63,"props":373,"children":375},{"className":374},[],[376],{"type":49,"value":377},"\"indexed_linear\"",{"type":43,"tag":304,"props":379,"children":380},{},[381],{"type":49,"value":339},{"type":43,"tag":304,"props":383,"children":384},{},[385],{"type":49,"value":386},"GPU, linear with indexed weights, sorted indices",{"type":43,"tag":52,"props":388,"children":390},{"id":389},"core-primitive-segmentedpolynomial",[391],{"type":49,"value":392},"Core primitive: SegmentedPolynomial",{"type":43,"tag":394,"props":395,"children":400},"pre",{"className":396,"code":397,"language":398,"meta":399,"style":399},"language-python shiki shiki-themes material-theme-lighter material-theme material-theme-palenight","import torch\nimport cuequivariance as cue\nimport cuequivariance_torch as cuet\n\n# Build a descriptor\ne = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2])\npoly = e.polynomial\n\n# Create the module\nsp = cuet.SegmentedPolynomial(poly, method=\"uniform_1d\")\n\n# Forward pass\nx = torch.randn(batch, 3, device=\"cuda\")\n[output] = sp([x])\n# output.shape == (batch, 9)  -- 1 + 3 + 5\n","python","",[401],{"type":43,"tag":63,"props":402,"children":403},{"__ignoreMap":399},[404,415,424,433,443,452,461,470,478,487,496,504,513,522,531],{"type":43,"tag":405,"props":406,"children":409},"span",{"class":407,"line":408},"line",1,[410],{"type":43,"tag":405,"props":411,"children":412},{},[413],{"type":49,"value":414},"import torch\n",{"type":43,"tag":405,"props":416,"children":418},{"class":407,"line":417},2,[419],{"type":43,"tag":405,"props":420,"children":421},{},[422],{"type":49,"value":423},"import cuequivariance as cue\n",{"type":43,"tag":405,"props":425,"children":427},{"class":407,"line":426},3,[428],{"type":43,"tag":405,"props":429,"children":430},{},[431],{"type":49,"value":432},"import cuequivariance_torch as cuet\n",{"type":43,"tag":405,"props":434,"children":436},{"class":407,"line":435},4,[437],{"type":43,"tag":405,"props":438,"children":440},{"emptyLinePlaceholder":439},true,[441],{"type":49,"value":442},"\n",{"type":43,"tag":405,"props":444,"children":446},{"class":407,"line":445},5,[447],{"type":43,"tag":405,"props":448,"children":449},{},[450],{"type":49,"value":451},"# Build a descriptor\n",{"type":43,"tag":405,"props":453,"children":455},{"class":407,"line":454},6,[456],{"type":43,"tag":405,"props":457,"children":458},{},[459],{"type":49,"value":460},"e = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2])\n",{"type":43,"tag":405,"props":462,"children":464},{"class":407,"line":463},7,[465],{"type":43,"tag":405,"props":466,"children":467},{},[468],{"type":49,"value":469},"poly = e.polynomial\n",{"type":43,"tag":405,"props":471,"children":473},{"class":407,"line":472},8,[474],{"type":43,"tag":405,"props":475,"children":476},{"emptyLinePlaceholder":439},[477],{"type":49,"value":442},{"type":43,"tag":405,"props":479,"children":481},{"class":407,"line":480},9,[482],{"type":43,"tag":405,"props":483,"children":484},{},[485],{"type":49,"value":486},"# Create the module\n",{"type":43,"tag":405,"props":488,"children":490},{"class":407,"line":489},10,[491],{"type":43,"tag":405,"props":492,"children":493},{},[494],{"type":49,"value":495},"sp = cuet.SegmentedPolynomial(poly, method=\"uniform_1d\")\n",{"type":43,"tag":405,"props":497,"children":499},{"class":407,"line":498},11,[500],{"type":43,"tag":405,"props":501,"children":502},{"emptyLinePlaceholder":439},[503],{"type":49,"value":442},{"type":43,"tag":405,"props":505,"children":507},{"class":407,"line":506},12,[508],{"type":43,"tag":405,"props":509,"children":510},{},[511],{"type":49,"value":512},"# Forward pass\n",{"type":43,"tag":405,"props":514,"children":516},{"class":407,"line":515},13,[517],{"type":43,"tag":405,"props":518,"children":519},{},[520],{"type":49,"value":521},"x = torch.randn(batch, 3, device=\"cuda\")\n",{"type":43,"tag":405,"props":523,"children":525},{"class":407,"line":524},14,[526],{"type":43,"tag":405,"props":527,"children":528},{},[529],{"type":49,"value":530},"[output] = sp([x])\n",{"type":43,"tag":405,"props":532,"children":534},{"class":407,"line":533},15,[535],{"type":43,"tag":405,"props":536,"children":537},{},[538],{"type":49,"value":539},"# output.shape == (batch, 9)  -- 1 + 3 + 5\n",{"type":43,"tag":541,"props":542,"children":544},"h3",{"id":543},"inputs-indexing-and-scatter",[545],{"type":49,"value":546},"Inputs, indexing, and scatter",{"type":43,"tag":394,"props":548,"children":550},{"className":396,"code":549,"language":398,"meta":399,"style":399},"e = cue.descriptors.channelwise_tensor_product(\n    16 * cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n)\npoly = e.polynomial\n\nsp = cuet.SegmentedPolynomial(poly, method=\"uniform_1d\")\n\nw = torch.randn(1, poly.inputs[0].size, device=\"cuda\")            # shared weights\nx1 = torch.randn(batch, poly.inputs[1].size, device=\"cuda\")       # batched input 1\nx2 = torch.randn(batch, poly.inputs[2].size, device=\"cuda\")       # batched input 2\n\n# Basic forward\n[out] = sp([w, x1, x2])\n\n# With input gathering (e.g., gather x1 by node index)\nsenders = torch.randint(0, num_nodes, (num_edges,), device=\"cuda\")\n[out] = sp([w, x1, x2], input_indices={1: senders})\n\n# With output scattering (accumulate into target nodes)\nreceivers = torch.randint(0, num_nodes, (num_edges,), device=\"cuda\")\n[out] = sp(\n    [w, x1, x2],\n    input_indices={1: senders},\n    output_indices={0: receivers},\n    output_shapes={0: torch.empty(num_nodes, 1, device=\"cuda\")},\n)\n",[551],{"type":43,"tag":63,"props":552,"children":553},{"__ignoreMap":399},[554,562,570,578,585,593,600,607,614,621,629,637,645,652,660,668,676,685,694,703,711,720,729,738,747,756,765,774],{"type":43,"tag":405,"props":555,"children":556},{"class":407,"line":408},[557],{"type":43,"tag":405,"props":558,"children":559},{},[560],{"type":49,"value":561},"e = cue.descriptors.channelwise_tensor_product(\n",{"type":43,"tag":405,"props":563,"children":564},{"class":407,"line":417},[565],{"type":43,"tag":405,"props":566,"children":567},{},[568],{"type":49,"value":569},"    16 * cue.Irreps(\"SO3\", \"0 + 1\"),\n",{"type":43,"tag":405,"props":571,"children":572},{"class":407,"line":426},[573],{"type":43,"tag":405,"props":574,"children":575},{},[576],{"type":49,"value":577},"    cue.Irreps(\"SO3\", \"0 + 1\"),\n",{"type":43,"tag":405,"props":579,"children":580},{"class":407,"line":435},[581],{"type":43,"tag":405,"props":582,"children":583},{},[584],{"type":49,"value":577},{"type":43,"tag":405,"props":586,"children":587},{"class":407,"line":445},[588],{"type":43,"tag":405,"props":589,"children":590},{},[591],{"type":49,"value":592},")\n",{"type":43,"tag":405,"props":594,"children":595},{"class":407,"line":454},[596],{"type":43,"tag":405,"props":597,"children":598},{},[599],{"type":49,"value":469},{"type":43,"tag":405,"props":601,"children":602},{"class":407,"line":463},[603],{"type":43,"tag":405,"props":604,"children":605},{"emptyLinePlaceholder":439},[606],{"type":49,"value":442},{"type":43,"tag":405,"props":608,"children":609},{"class":407,"line":472},[610],{"type":43,"tag":405,"props":611,"children":612},{},[613],{"type":49,"value":495},{"type":43,"tag":405,"props":615,"children":616},{"class":407,"line":480},[617],{"type":43,"tag":405,"props":618,"children":619},{"emptyLinePlaceholder":439},[620],{"type":49,"value":442},{"type":43,"tag":405,"props":622,"children":623},{"class":407,"line":489},[624],{"type":43,"tag":405,"props":625,"children":626},{},[627],{"type":49,"value":628},"w = torch.randn(1, poly.inputs[0].size, device=\"cuda\")            # shared weights\n",{"type":43,"tag":405,"props":630,"children":631},{"class":407,"line":498},[632],{"type":43,"tag":405,"props":633,"children":634},{},[635],{"type":49,"value":636},"x1 = torch.randn(batch, poly.inputs[1].size, device=\"cuda\")       # batched input 1\n",{"type":43,"tag":405,"props":638,"children":639},{"class":407,"line":506},[640],{"type":43,"tag":405,"props":641,"children":642},{},[643],{"type":49,"value":644},"x2 = torch.randn(batch, poly.inputs[2].size, device=\"cuda\")       # batched input 2\n",{"type":43,"tag":405,"props":646,"children":647},{"class":407,"line":515},[648],{"type":43,"tag":405,"props":649,"children":650},{"emptyLinePlaceholder":439},[651],{"type":49,"value":442},{"type":43,"tag":405,"props":653,"children":654},{"class":407,"line":524},[655],{"type":43,"tag":405,"props":656,"children":657},{},[658],{"type":49,"value":659},"# Basic forward\n",{"type":43,"tag":405,"props":661,"children":662},{"class":407,"line":533},[663],{"type":43,"tag":405,"props":664,"children":665},{},[666],{"type":49,"value":667},"[out] = sp([w, x1, x2])\n",{"type":43,"tag":405,"props":669,"children":671},{"class":407,"line":670},16,[672],{"type":43,"tag":405,"props":673,"children":674},{"emptyLinePlaceholder":439},[675],{"type":49,"value":442},{"type":43,"tag":405,"props":677,"children":679},{"class":407,"line":678},17,[680],{"type":43,"tag":405,"props":681,"children":682},{},[683],{"type":49,"value":684},"# With input gathering (e.g., gather x1 by node index)\n",{"type":43,"tag":405,"props":686,"children":688},{"class":407,"line":687},18,[689],{"type":43,"tag":405,"props":690,"children":691},{},[692],{"type":49,"value":693},"senders = torch.randint(0, num_nodes, (num_edges,), device=\"cuda\")\n",{"type":43,"tag":405,"props":695,"children":697},{"class":407,"line":696},19,[698],{"type":43,"tag":405,"props":699,"children":700},{},[701],{"type":49,"value":702},"[out] = sp([w, x1, x2], input_indices={1: senders})\n",{"type":43,"tag":405,"props":704,"children":706},{"class":407,"line":705},20,[707],{"type":43,"tag":405,"props":708,"children":709},{"emptyLinePlaceholder":439},[710],{"type":49,"value":442},{"type":43,"tag":405,"props":712,"children":714},{"class":407,"line":713},21,[715],{"type":43,"tag":405,"props":716,"children":717},{},[718],{"type":49,"value":719},"# With output scattering (accumulate into target nodes)\n",{"type":43,"tag":405,"props":721,"children":723},{"class":407,"line":722},22,[724],{"type":43,"tag":405,"props":725,"children":726},{},[727],{"type":49,"value":728},"receivers = torch.randint(0, num_nodes, (num_edges,), device=\"cuda\")\n",{"type":43,"tag":405,"props":730,"children":732},{"class":407,"line":731},23,[733],{"type":43,"tag":405,"props":734,"children":735},{},[736],{"type":49,"value":737},"[out] = sp(\n",{"type":43,"tag":405,"props":739,"children":741},{"class":407,"line":740},24,[742],{"type":43,"tag":405,"props":743,"children":744},{},[745],{"type":49,"value":746},"    [w, x1, x2],\n",{"type":43,"tag":405,"props":748,"children":750},{"class":407,"line":749},25,[751],{"type":43,"tag":405,"props":752,"children":753},{},[754],{"type":49,"value":755},"    input_indices={1: senders},\n",{"type":43,"tag":405,"props":757,"children":759},{"class":407,"line":758},26,[760],{"type":43,"tag":405,"props":761,"children":762},{},[763],{"type":49,"value":764},"    output_indices={0: receivers},\n",{"type":43,"tag":405,"props":766,"children":768},{"class":407,"line":767},27,[769],{"type":43,"tag":405,"props":770,"children":771},{},[772],{"type":49,"value":773},"    output_shapes={0: torch.empty(num_nodes, 1, device=\"cuda\")},\n",{"type":43,"tag":405,"props":775,"children":777},{"class":407,"line":776},28,[778],{"type":43,"tag":405,"props":779,"children":780},{},[781],{"type":49,"value":592},{"type":43,"tag":541,"props":783,"children":785},{"id":784},"math-dtype-control",[786],{"type":49,"value":787},"Math dtype control",{"type":43,"tag":394,"props":789,"children":791},{"className":396,"code":790,"language":398,"meta":399,"style":399},"# Compute in float32 regardless of input dtype\nsp = cuet.SegmentedPolynomial(poly, method=\"fused_tp\", math_dtype=torch.float32)\n\n# For fused_tp, math_dtype must be float32 or float64\n# For naive, any torch.dtype works\n# For uniform_1d, float32 or float64 (auto-selects float32 if input is e.g. float16)\n",[792],{"type":43,"tag":63,"props":793,"children":794},{"__ignoreMap":399},[795,803,811,818,826,834],{"type":43,"tag":405,"props":796,"children":797},{"class":407,"line":408},[798],{"type":43,"tag":405,"props":799,"children":800},{},[801],{"type":49,"value":802},"# Compute in float32 regardless of input dtype\n",{"type":43,"tag":405,"props":804,"children":805},{"class":407,"line":417},[806],{"type":43,"tag":405,"props":807,"children":808},{},[809],{"type":49,"value":810},"sp = cuet.SegmentedPolynomial(poly, method=\"fused_tp\", math_dtype=torch.float32)\n",{"type":43,"tag":405,"props":812,"children":813},{"class":407,"line":426},[814],{"type":43,"tag":405,"props":815,"children":816},{"emptyLinePlaceholder":439},[817],{"type":49,"value":442},{"type":43,"tag":405,"props":819,"children":820},{"class":407,"line":435},[821],{"type":43,"tag":405,"props":822,"children":823},{},[824],{"type":49,"value":825},"# For fused_tp, math_dtype must be float32 or float64\n",{"type":43,"tag":405,"props":827,"children":828},{"class":407,"line":445},[829],{"type":43,"tag":405,"props":830,"children":831},{},[832],{"type":49,"value":833},"# For naive, any torch.dtype works\n",{"type":43,"tag":405,"props":835,"children":836},{"class":407,"line":454},[837],{"type":43,"tag":405,"props":838,"children":839},{},[840],{"type":49,"value":841},"# For uniform_1d, float32 or float64 (auto-selects float32 if input is e.g. float16)\n",{"type":43,"tag":52,"props":843,"children":845},{"id":844},"high-level-operations",[846],{"type":49,"value":126},{"type":43,"tag":59,"props":848,"children":849},{},[850,852,857,859,865],{"type":49,"value":851},"All operations are ",{"type":43,"tag":63,"props":853,"children":855},{"className":854},[],[856],{"type":49,"value":116},{"type":49,"value":858}," subclasses. They wrap ",{"type":43,"tag":63,"props":860,"children":862},{"className":861},[],[863],{"type":49,"value":864},"SegmentedPolynomial",{"type":49,"value":866}," and handle layout transposition automatically.",{"type":43,"tag":541,"props":868,"children":870},{"id":869},"memory-layout",[871],{"type":49,"value":872},"Memory layout",{"type":43,"tag":59,"props":874,"children":875},{},[876,882,884,890],{"type":43,"tag":63,"props":877,"children":879},{"className":878},[],[880],{"type":49,"value":881},"IrrepsLayout",{"type":49,"value":883}," controls memory order within each ",{"type":43,"tag":63,"props":885,"children":887},{"className":886},[],[888],{"type":49,"value":889},"(mul, ir)",{"type":49,"value":891}," block:",{"type":43,"tag":893,"props":894,"children":895},"ul",{},[896,919],{"type":43,"tag":92,"props":897,"children":898},{},[899,905,907,913,914],{"type":43,"tag":63,"props":900,"children":902},{"className":901},[],[903],{"type":49,"value":904},"cue.mul_ir",{"type":49,"value":906},": data ordered as ",{"type":43,"tag":63,"props":908,"children":910},{"className":909},[],[911],{"type":49,"value":912},"(mul, ir.dim)",{"type":49,"value":110},{"type":43,"tag":96,"props":915,"children":916},{},[917],{"type":49,"value":918},"default, compatible with e3nn",{"type":43,"tag":92,"props":920,"children":921},{},[922,928,929,935,936],{"type":43,"tag":63,"props":923,"children":925},{"className":924},[],[926],{"type":49,"value":927},"cue.ir_mul",{"type":49,"value":906},{"type":43,"tag":63,"props":930,"children":932},{"className":931},[],[933],{"type":49,"value":934},"(ir.dim, mul)",{"type":49,"value":110},{"type":43,"tag":96,"props":937,"children":938},{},[939],{"type":49,"value":940},"used internally by descriptors",{"type":43,"tag":59,"props":942,"children":943},{},[944,946,952,954,960,961,967,968,974],{"type":49,"value":945},"Operations accept ",{"type":43,"tag":63,"props":947,"children":949},{"className":948},[],[950],{"type":49,"value":951},"layout",{"type":49,"value":953}," (applies to all), or per-operand ",{"type":43,"tag":63,"props":955,"children":957},{"className":956},[],[958],{"type":49,"value":959},"layout_in1",{"type":49,"value":143},{"type":43,"tag":63,"props":962,"children":964},{"className":963},[],[965],{"type":49,"value":966},"layout_in2",{"type":49,"value":143},{"type":43,"tag":63,"props":969,"children":971},{"className":970},[],[972],{"type":49,"value":973},"layout_out",{"type":49,"value":975},".",{"type":43,"tag":541,"props":977,"children":979},{"id":978},"channelwisetensorproduct",[980],{"type":49,"value":141},{"type":43,"tag":59,"props":982,"children":983},{},[984,986,992,994,1000],{"type":49,"value":985},"Channel-wise tensor product: pairs channels of ",{"type":43,"tag":63,"props":987,"children":989},{"className":988},[],[990],{"type":49,"value":991},"x1",{"type":49,"value":993}," with channels of ",{"type":43,"tag":63,"props":995,"children":997},{"className":996},[],[998],{"type":49,"value":999},"x2",{"type":49,"value":975},{"type":43,"tag":394,"props":1002,"children":1004},{"className":396,"code":1003,"language":398,"meta":399,"style":399},"# With internal weights (default: shared_weights=True, internal_weights=True)\ntp = cuet.ChannelWiseTensorProduct(\n    cue.Irreps(\"SO3\", \"32x0 + 32x1\"),   # irreps_in1\n    cue.Irreps(\"SO3\", \"0 + 1\"),          # irreps_in2\n    layout=cue.mul_ir,\n    device=\"cuda\",\n    dtype=torch.float32,\n)\n# tp.weight_numel -- number of weight parameters\n# tp.irreps_out -- output irreps (auto-computed)\n\nx1 = torch.randn(batch, tp.irreps_in1.dim, device=\"cuda\")\nx2 = torch.randn(batch, tp.irreps_in2.dim, device=\"cuda\")\n\nout = tp(x1, x2)  # uses internal weight parameter\n# out.shape == (batch, tp.irreps_out.dim)\n\n# With external weights (shared_weights=False)\ntp = cuet.ChannelWiseTensorProduct(\n    cue.Irreps(\"SO3\", \"32x0 + 32x1\"),\n    cue.Irreps(\"SO3\", \"0 + 1\"),\n    layout=cue.mul_ir,\n    shared_weights=False,\n    device=\"cuda\",\n)\nw = torch.randn(batch, tp.weight_numel, device=\"cuda\")\nout = tp(x1, x2, weight=w)\n\n# With gather\u002Fscatter for graph neural networks\nout = tp(x1, x2, weight=w, indices_1=senders, indices_out=receivers, size_out=num_nodes)\n",[1005],{"type":43,"tag":63,"props":1006,"children":1007},{"__ignoreMap":399},[1008,1016,1024,1032,1040,1048,1056,1064,1071,1079,1087,1094,1102,1110,1117,1125,1133,1140,1148,1155,1163,1170,1177,1185,1192,1199,1207,1215,1222,1231],{"type":43,"tag":405,"props":1009,"children":1010},{"class":407,"line":408},[1011],{"type":43,"tag":405,"props":1012,"children":1013},{},[1014],{"type":49,"value":1015},"# With internal weights (default: shared_weights=True, internal_weights=True)\n",{"type":43,"tag":405,"props":1017,"children":1018},{"class":407,"line":417},[1019],{"type":43,"tag":405,"props":1020,"children":1021},{},[1022],{"type":49,"value":1023},"tp = cuet.ChannelWiseTensorProduct(\n",{"type":43,"tag":405,"props":1025,"children":1026},{"class":407,"line":426},[1027],{"type":43,"tag":405,"props":1028,"children":1029},{},[1030],{"type":49,"value":1031},"    cue.Irreps(\"SO3\", \"32x0 + 32x1\"),   # irreps_in1\n",{"type":43,"tag":405,"props":1033,"children":1034},{"class":407,"line":435},[1035],{"type":43,"tag":405,"props":1036,"children":1037},{},[1038],{"type":49,"value":1039},"    cue.Irreps(\"SO3\", \"0 + 1\"),          # irreps_in2\n",{"type":43,"tag":405,"props":1041,"children":1042},{"class":407,"line":445},[1043],{"type":43,"tag":405,"props":1044,"children":1045},{},[1046],{"type":49,"value":1047},"    layout=cue.mul_ir,\n",{"type":43,"tag":405,"props":1049,"children":1050},{"class":407,"line":454},[1051],{"type":43,"tag":405,"props":1052,"children":1053},{},[1054],{"type":49,"value":1055},"    device=\"cuda\",\n",{"type":43,"tag":405,"props":1057,"children":1058},{"class":407,"line":463},[1059],{"type":43,"tag":405,"props":1060,"children":1061},{},[1062],{"type":49,"value":1063},"    dtype=torch.float32,\n",{"type":43,"tag":405,"props":1065,"children":1066},{"class":407,"line":472},[1067],{"type":43,"tag":405,"props":1068,"children":1069},{},[1070],{"type":49,"value":592},{"type":43,"tag":405,"props":1072,"children":1073},{"class":407,"line":480},[1074],{"type":43,"tag":405,"props":1075,"children":1076},{},[1077],{"type":49,"value":1078},"# tp.weight_numel -- number of weight parameters\n",{"type":43,"tag":405,"props":1080,"children":1081},{"class":407,"line":489},[1082],{"type":43,"tag":405,"props":1083,"children":1084},{},[1085],{"type":49,"value":1086},"# tp.irreps_out -- output irreps (auto-computed)\n",{"type":43,"tag":405,"props":1088,"children":1089},{"class":407,"line":498},[1090],{"type":43,"tag":405,"props":1091,"children":1092},{"emptyLinePlaceholder":439},[1093],{"type":49,"value":442},{"type":43,"tag":405,"props":1095,"children":1096},{"class":407,"line":506},[1097],{"type":43,"tag":405,"props":1098,"children":1099},{},[1100],{"type":49,"value":1101},"x1 = torch.randn(batch, tp.irreps_in1.dim, device=\"cuda\")\n",{"type":43,"tag":405,"props":1103,"children":1104},{"class":407,"line":515},[1105],{"type":43,"tag":405,"props":1106,"children":1107},{},[1108],{"type":49,"value":1109},"x2 = torch.randn(batch, tp.irreps_in2.dim, device=\"cuda\")\n",{"type":43,"tag":405,"props":1111,"children":1112},{"class":407,"line":524},[1113],{"type":43,"tag":405,"props":1114,"children":1115},{"emptyLinePlaceholder":439},[1116],{"type":49,"value":442},{"type":43,"tag":405,"props":1118,"children":1119},{"class":407,"line":533},[1120],{"type":43,"tag":405,"props":1121,"children":1122},{},[1123],{"type":49,"value":1124},"out = tp(x1, x2)  # uses internal weight parameter\n",{"type":43,"tag":405,"props":1126,"children":1127},{"class":407,"line":670},[1128],{"type":43,"tag":405,"props":1129,"children":1130},{},[1131],{"type":49,"value":1132},"# out.shape == (batch, tp.irreps_out.dim)\n",{"type":43,"tag":405,"props":1134,"children":1135},{"class":407,"line":678},[1136],{"type":43,"tag":405,"props":1137,"children":1138},{"emptyLinePlaceholder":439},[1139],{"type":49,"value":442},{"type":43,"tag":405,"props":1141,"children":1142},{"class":407,"line":687},[1143],{"type":43,"tag":405,"props":1144,"children":1145},{},[1146],{"type":49,"value":1147},"# With external weights (shared_weights=False)\n",{"type":43,"tag":405,"props":1149,"children":1150},{"class":407,"line":696},[1151],{"type":43,"tag":405,"props":1152,"children":1153},{},[1154],{"type":49,"value":1023},{"type":43,"tag":405,"props":1156,"children":1157},{"class":407,"line":705},[1158],{"type":43,"tag":405,"props":1159,"children":1160},{},[1161],{"type":49,"value":1162},"    cue.Irreps(\"SO3\", \"32x0 + 32x1\"),\n",{"type":43,"tag":405,"props":1164,"children":1165},{"class":407,"line":713},[1166],{"type":43,"tag":405,"props":1167,"children":1168},{},[1169],{"type":49,"value":577},{"type":43,"tag":405,"props":1171,"children":1172},{"class":407,"line":722},[1173],{"type":43,"tag":405,"props":1174,"children":1175},{},[1176],{"type":49,"value":1047},{"type":43,"tag":405,"props":1178,"children":1179},{"class":407,"line":731},[1180],{"type":43,"tag":405,"props":1181,"children":1182},{},[1183],{"type":49,"value":1184},"    shared_weights=False,\n",{"type":43,"tag":405,"props":1186,"children":1187},{"class":407,"line":740},[1188],{"type":43,"tag":405,"props":1189,"children":1190},{},[1191],{"type":49,"value":1055},{"type":43,"tag":405,"props":1193,"children":1194},{"class":407,"line":749},[1195],{"type":43,"tag":405,"props":1196,"children":1197},{},[1198],{"type":49,"value":592},{"type":43,"tag":405,"props":1200,"children":1201},{"class":407,"line":758},[1202],{"type":43,"tag":405,"props":1203,"children":1204},{},[1205],{"type":49,"value":1206},"w = torch.randn(batch, tp.weight_numel, device=\"cuda\")\n",{"type":43,"tag":405,"props":1208,"children":1209},{"class":407,"line":767},[1210],{"type":43,"tag":405,"props":1211,"children":1212},{},[1213],{"type":49,"value":1214},"out = tp(x1, x2, weight=w)\n",{"type":43,"tag":405,"props":1216,"children":1217},{"class":407,"line":776},[1218],{"type":43,"tag":405,"props":1219,"children":1220},{"emptyLinePlaceholder":439},[1221],{"type":49,"value":442},{"type":43,"tag":405,"props":1223,"children":1225},{"class":407,"line":1224},29,[1226],{"type":43,"tag":405,"props":1227,"children":1228},{},[1229],{"type":49,"value":1230},"# With gather\u002Fscatter for graph neural networks\n",{"type":43,"tag":405,"props":1232,"children":1234},{"class":407,"line":1233},30,[1235],{"type":43,"tag":405,"props":1236,"children":1237},{},[1238],{"type":49,"value":1239},"out = tp(x1, x2, weight=w, indices_1=senders, indices_out=receivers, size_out=num_nodes)\n",{"type":43,"tag":59,"props":1241,"children":1242},{},[1243,1245,1250,1252,1257],{"type":49,"value":1244},"Default method: ",{"type":43,"tag":63,"props":1246,"children":1248},{"className":1247},[],[1249],{"type":49,"value":334},{"type":49,"value":1251}," if segments are uniform, else ",{"type":43,"tag":63,"props":1253,"children":1255},{"className":1254},[],[1256],{"type":49,"value":312},{"type":49,"value":975},{"type":43,"tag":541,"props":1259,"children":1261},{"id":1260},"fullyconnectedtensorproduct",[1262],{"type":49,"value":149},{"type":43,"tag":59,"props":1264,"children":1265},{},[1266],{"type":49,"value":1267},"All input irrep pairs contribute to all output irreps (dense contraction).",{"type":43,"tag":394,"props":1269,"children":1271},{"className":396,"code":1270,"language":398,"meta":399,"style":399},"tp = cuet.FullyConnectedTensorProduct(\n    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),    # irreps_in1\n    cue.Irreps(\"O3\", \"0e + 1o\"),         # irreps_in2\n    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),    # irreps_out\n    layout=cue.mul_ir,\n    internal_weights=True,\n    device=\"cuda\",\n)\n\nout = tp(x1, x2)  # uses internal weights\n# or: out = tp(x1, x2, weight=w)  # external weights\n",[1272],{"type":43,"tag":63,"props":1273,"children":1274},{"__ignoreMap":399},[1275,1283,1291,1299,1307,1314,1322,1329,1336,1343,1351],{"type":43,"tag":405,"props":1276,"children":1277},{"class":407,"line":408},[1278],{"type":43,"tag":405,"props":1279,"children":1280},{},[1281],{"type":49,"value":1282},"tp = cuet.FullyConnectedTensorProduct(\n",{"type":43,"tag":405,"props":1284,"children":1285},{"class":407,"line":417},[1286],{"type":43,"tag":405,"props":1287,"children":1288},{},[1289],{"type":49,"value":1290},"    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),    # irreps_in1\n",{"type":43,"tag":405,"props":1292,"children":1293},{"class":407,"line":426},[1294],{"type":43,"tag":405,"props":1295,"children":1296},{},[1297],{"type":49,"value":1298},"    cue.Irreps(\"O3\", \"0e + 1o\"),         # irreps_in2\n",{"type":43,"tag":405,"props":1300,"children":1301},{"class":407,"line":435},[1302],{"type":43,"tag":405,"props":1303,"children":1304},{},[1305],{"type":49,"value":1306},"    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),    # irreps_out\n",{"type":43,"tag":405,"props":1308,"children":1309},{"class":407,"line":445},[1310],{"type":43,"tag":405,"props":1311,"children":1312},{},[1313],{"type":49,"value":1047},{"type":43,"tag":405,"props":1315,"children":1316},{"class":407,"line":454},[1317],{"type":43,"tag":405,"props":1318,"children":1319},{},[1320],{"type":49,"value":1321},"    internal_weights=True,\n",{"type":43,"tag":405,"props":1323,"children":1324},{"class":407,"line":463},[1325],{"type":43,"tag":405,"props":1326,"children":1327},{},[1328],{"type":49,"value":1055},{"type":43,"tag":405,"props":1330,"children":1331},{"class":407,"line":472},[1332],{"type":43,"tag":405,"props":1333,"children":1334},{},[1335],{"type":49,"value":592},{"type":43,"tag":405,"props":1337,"children":1338},{"class":407,"line":480},[1339],{"type":43,"tag":405,"props":1340,"children":1341},{"emptyLinePlaceholder":439},[1342],{"type":49,"value":442},{"type":43,"tag":405,"props":1344,"children":1345},{"class":407,"line":489},[1346],{"type":43,"tag":405,"props":1347,"children":1348},{},[1349],{"type":49,"value":1350},"out = tp(x1, x2)  # uses internal weights\n",{"type":43,"tag":405,"props":1352,"children":1353},{"class":407,"line":498},[1354],{"type":43,"tag":405,"props":1355,"children":1356},{},[1357],{"type":49,"value":1358},"# or: out = tp(x1, x2, weight=w)  # external weights\n",{"type":43,"tag":59,"props":1360,"children":1361},{},[1362,1363,1368],{"type":49,"value":1244},{"type":43,"tag":63,"props":1364,"children":1366},{"className":1365},[],[1367],{"type":49,"value":356},{"type":49,"value":975},{"type":43,"tag":541,"props":1370,"children":1372},{"id":1371},"linear",[1373],{"type":49,"value":156},{"type":43,"tag":59,"props":1375,"children":1376},{},[1377],{"type":49,"value":1378},"Equivariant linear layer (weight-only, no second input).",{"type":43,"tag":394,"props":1380,"children":1382},{"className":396,"code":1381,"language":398,"meta":399,"style":399},"linear = cuet.Linear(\n    cue.Irreps(\"SO3\", \"4x0 + 2x1\"),     # irreps_in\n    cue.Irreps(\"SO3\", \"3x0 + 5x1\"),     # irreps_out\n    layout=cue.mul_ir,\n    internal_weights=True,\n    device=\"cuda\",\n)\n\nout = linear(x)\n\n# Species-indexed weights (different weights per atom type)\nlinear = cuet.Linear(\n    irreps_in, irreps_out,\n    weight_classes=50,   # 50 different weight sets\n    internal_weights=True,\n    device=\"cuda\",\n)\nout = linear(x, weight_indices=species_indices)  # species_indices: (batch,) int tensor\n",[1383],{"type":43,"tag":63,"props":1384,"children":1385},{"__ignoreMap":399},[1386,1394,1402,1410,1417,1424,1431,1438,1445,1453,1460,1468,1475,1483,1491,1498,1505,1512],{"type":43,"tag":405,"props":1387,"children":1388},{"class":407,"line":408},[1389],{"type":43,"tag":405,"props":1390,"children":1391},{},[1392],{"type":49,"value":1393},"linear = cuet.Linear(\n",{"type":43,"tag":405,"props":1395,"children":1396},{"class":407,"line":417},[1397],{"type":43,"tag":405,"props":1398,"children":1399},{},[1400],{"type":49,"value":1401},"    cue.Irreps(\"SO3\", \"4x0 + 2x1\"),     # irreps_in\n",{"type":43,"tag":405,"props":1403,"children":1404},{"class":407,"line":426},[1405],{"type":43,"tag":405,"props":1406,"children":1407},{},[1408],{"type":49,"value":1409},"    cue.Irreps(\"SO3\", \"3x0 + 5x1\"),     # irreps_out\n",{"type":43,"tag":405,"props":1411,"children":1412},{"class":407,"line":435},[1413],{"type":43,"tag":405,"props":1414,"children":1415},{},[1416],{"type":49,"value":1047},{"type":43,"tag":405,"props":1418,"children":1419},{"class":407,"line":445},[1420],{"type":43,"tag":405,"props":1421,"children":1422},{},[1423],{"type":49,"value":1321},{"type":43,"tag":405,"props":1425,"children":1426},{"class":407,"line":454},[1427],{"type":43,"tag":405,"props":1428,"children":1429},{},[1430],{"type":49,"value":1055},{"type":43,"tag":405,"props":1432,"children":1433},{"class":407,"line":463},[1434],{"type":43,"tag":405,"props":1435,"children":1436},{},[1437],{"type":49,"value":592},{"type":43,"tag":405,"props":1439,"children":1440},{"class":407,"line":472},[1441],{"type":43,"tag":405,"props":1442,"children":1443},{"emptyLinePlaceholder":439},[1444],{"type":49,"value":442},{"type":43,"tag":405,"props":1446,"children":1447},{"class":407,"line":480},[1448],{"type":43,"tag":405,"props":1449,"children":1450},{},[1451],{"type":49,"value":1452},"out = linear(x)\n",{"type":43,"tag":405,"props":1454,"children":1455},{"class":407,"line":489},[1456],{"type":43,"tag":405,"props":1457,"children":1458},{"emptyLinePlaceholder":439},[1459],{"type":49,"value":442},{"type":43,"tag":405,"props":1461,"children":1462},{"class":407,"line":498},[1463],{"type":43,"tag":405,"props":1464,"children":1465},{},[1466],{"type":49,"value":1467},"# Species-indexed weights (different weights per atom type)\n",{"type":43,"tag":405,"props":1469,"children":1470},{"class":407,"line":506},[1471],{"type":43,"tag":405,"props":1472,"children":1473},{},[1474],{"type":49,"value":1393},{"type":43,"tag":405,"props":1476,"children":1477},{"class":407,"line":515},[1478],{"type":43,"tag":405,"props":1479,"children":1480},{},[1481],{"type":49,"value":1482},"    irreps_in, irreps_out,\n",{"type":43,"tag":405,"props":1484,"children":1485},{"class":407,"line":524},[1486],{"type":43,"tag":405,"props":1487,"children":1488},{},[1489],{"type":49,"value":1490},"    weight_classes=50,   # 50 different weight sets\n",{"type":43,"tag":405,"props":1492,"children":1493},{"class":407,"line":533},[1494],{"type":43,"tag":405,"props":1495,"children":1496},{},[1497],{"type":49,"value":1321},{"type":43,"tag":405,"props":1499,"children":1500},{"class":407,"line":670},[1501],{"type":43,"tag":405,"props":1502,"children":1503},{},[1504],{"type":49,"value":1055},{"type":43,"tag":405,"props":1506,"children":1507},{"class":407,"line":678},[1508],{"type":43,"tag":405,"props":1509,"children":1510},{},[1511],{"type":49,"value":592},{"type":43,"tag":405,"props":1513,"children":1514},{"class":407,"line":687},[1515],{"type":43,"tag":405,"props":1516,"children":1517},{},[1518],{"type":49,"value":1519},"out = linear(x, weight_indices=species_indices)  # species_indices: (batch,) int tensor\n",{"type":43,"tag":59,"props":1521,"children":1522},{},[1523,1524,1529,1531,1537],{"type":49,"value":1244},{"type":43,"tag":63,"props":1525,"children":1527},{"className":1526},[],[1528],{"type":49,"value":312},{"type":49,"value":1530},". Use ",{"type":43,"tag":63,"props":1532,"children":1534},{"className":1533},[],[1535],{"type":49,"value":1536},"method=\"fused_tp\"",{"type":49,"value":1538}," for CUDA acceleration.",{"type":43,"tag":541,"props":1540,"children":1542},{"id":1541},"symmetriccontraction",[1543],{"type":49,"value":163},{"type":43,"tag":59,"props":1545,"children":1546},{},[1547],{"type":49,"value":1548},"MACE-style symmetric contraction with element-indexed weights.",{"type":43,"tag":394,"props":1550,"children":1552},{"className":396,"code":1551,"language":398,"meta":399,"style":399},"sc = cuet.SymmetricContraction(\n    cue.Irreps(\"O3\", \"32x0e + 32x1o\"),  # irreps_in (uniform mul required)\n    cue.Irreps(\"O3\", \"32x0e\"),           # irreps_out (uniform mul required)\n    contraction_degree=3,\n    num_elements=95,                      # number of chemical elements\n    layout=cue.ir_mul,\n    dtype=torch.float32,\n    device=\"cuda\",\n)\n\n# indices: (batch,) int tensor selecting which element weights to use\nout = sc(x, indices)\n# out.shape == (batch, irreps_out.dim)\n",[1553],{"type":43,"tag":63,"props":1554,"children":1555},{"__ignoreMap":399},[1556,1564,1572,1580,1588,1596,1604,1611,1618,1625,1632,1640,1648],{"type":43,"tag":405,"props":1557,"children":1558},{"class":407,"line":408},[1559],{"type":43,"tag":405,"props":1560,"children":1561},{},[1562],{"type":49,"value":1563},"sc = cuet.SymmetricContraction(\n",{"type":43,"tag":405,"props":1565,"children":1566},{"class":407,"line":417},[1567],{"type":43,"tag":405,"props":1568,"children":1569},{},[1570],{"type":49,"value":1571},"    cue.Irreps(\"O3\", \"32x0e + 32x1o\"),  # irreps_in (uniform mul required)\n",{"type":43,"tag":405,"props":1573,"children":1574},{"class":407,"line":426},[1575],{"type":43,"tag":405,"props":1576,"children":1577},{},[1578],{"type":49,"value":1579},"    cue.Irreps(\"O3\", \"32x0e\"),           # irreps_out (uniform mul required)\n",{"type":43,"tag":405,"props":1581,"children":1582},{"class":407,"line":435},[1583],{"type":43,"tag":405,"props":1584,"children":1585},{},[1586],{"type":49,"value":1587},"    contraction_degree=3,\n",{"type":43,"tag":405,"props":1589,"children":1590},{"class":407,"line":445},[1591],{"type":43,"tag":405,"props":1592,"children":1593},{},[1594],{"type":49,"value":1595},"    num_elements=95,                      # number of chemical elements\n",{"type":43,"tag":405,"props":1597,"children":1598},{"class":407,"line":454},[1599],{"type":43,"tag":405,"props":1600,"children":1601},{},[1602],{"type":49,"value":1603},"    layout=cue.ir_mul,\n",{"type":43,"tag":405,"props":1605,"children":1606},{"class":407,"line":463},[1607],{"type":43,"tag":405,"props":1608,"children":1609},{},[1610],{"type":49,"value":1063},{"type":43,"tag":405,"props":1612,"children":1613},{"class":407,"line":472},[1614],{"type":43,"tag":405,"props":1615,"children":1616},{},[1617],{"type":49,"value":1055},{"type":43,"tag":405,"props":1619,"children":1620},{"class":407,"line":480},[1621],{"type":43,"tag":405,"props":1622,"children":1623},{},[1624],{"type":49,"value":592},{"type":43,"tag":405,"props":1626,"children":1627},{"class":407,"line":489},[1628],{"type":43,"tag":405,"props":1629,"children":1630},{"emptyLinePlaceholder":439},[1631],{"type":49,"value":442},{"type":43,"tag":405,"props":1633,"children":1634},{"class":407,"line":498},[1635],{"type":43,"tag":405,"props":1636,"children":1637},{},[1638],{"type":49,"value":1639},"# indices: (batch,) int tensor selecting which element weights to use\n",{"type":43,"tag":405,"props":1641,"children":1642},{"class":407,"line":506},[1643],{"type":43,"tag":405,"props":1644,"children":1645},{},[1646],{"type":49,"value":1647},"out = sc(x, indices)\n",{"type":43,"tag":405,"props":1649,"children":1650},{"class":407,"line":515},[1651],{"type":43,"tag":405,"props":1652,"children":1653},{},[1654],{"type":49,"value":1655},"# out.shape == (batch, irreps_out.dim)\n",{"type":43,"tag":59,"props":1657,"children":1658},{},[1659,1660,1665,1666,1671],{"type":49,"value":1244},{"type":43,"tag":63,"props":1661,"children":1663},{"className":1662},[],[1664],{"type":49,"value":334},{"type":49,"value":1251},{"type":43,"tag":63,"props":1667,"children":1669},{"className":1668},[],[1670],{"type":49,"value":312},{"type":49,"value":975},{"type":43,"tag":541,"props":1673,"children":1675},{"id":1674},"sphericalharmonics",[1676],{"type":49,"value":170},{"type":43,"tag":394,"props":1678,"children":1680},{"className":396,"code":1679,"language":398,"meta":399,"style":399},"sh = cuet.SphericalHarmonics(\n    ls=[0, 1, 2, 3],\n    normalize=True,\n    device=\"cuda\",\n)\n\nvectors = torch.randn(batch, 3, device=\"cuda\")\nout = sh(vectors)\n# out.shape == (batch, 1 + 3 + 5 + 7)  -- sum of 2l+1\n",[1681],{"type":43,"tag":63,"props":1682,"children":1683},{"__ignoreMap":399},[1684,1692,1700,1708,1715,1722,1729,1737,1745],{"type":43,"tag":405,"props":1685,"children":1686},{"class":407,"line":408},[1687],{"type":43,"tag":405,"props":1688,"children":1689},{},[1690],{"type":49,"value":1691},"sh = cuet.SphericalHarmonics(\n",{"type":43,"tag":405,"props":1693,"children":1694},{"class":407,"line":417},[1695],{"type":43,"tag":405,"props":1696,"children":1697},{},[1698],{"type":49,"value":1699},"    ls=[0, 1, 2, 3],\n",{"type":43,"tag":405,"props":1701,"children":1702},{"class":407,"line":426},[1703],{"type":43,"tag":405,"props":1704,"children":1705},{},[1706],{"type":49,"value":1707},"    normalize=True,\n",{"type":43,"tag":405,"props":1709,"children":1710},{"class":407,"line":435},[1711],{"type":43,"tag":405,"props":1712,"children":1713},{},[1714],{"type":49,"value":1055},{"type":43,"tag":405,"props":1716,"children":1717},{"class":407,"line":445},[1718],{"type":43,"tag":405,"props":1719,"children":1720},{},[1721],{"type":49,"value":592},{"type":43,"tag":405,"props":1723,"children":1724},{"class":407,"line":454},[1725],{"type":43,"tag":405,"props":1726,"children":1727},{"emptyLinePlaceholder":439},[1728],{"type":49,"value":442},{"type":43,"tag":405,"props":1730,"children":1731},{"class":407,"line":463},[1732],{"type":43,"tag":405,"props":1733,"children":1734},{},[1735],{"type":49,"value":1736},"vectors = torch.randn(batch, 3, device=\"cuda\")\n",{"type":43,"tag":405,"props":1738,"children":1739},{"class":407,"line":472},[1740],{"type":43,"tag":405,"props":1741,"children":1742},{},[1743],{"type":49,"value":1744},"out = sh(vectors)\n",{"type":43,"tag":405,"props":1746,"children":1747},{"class":407,"line":480},[1748],{"type":43,"tag":405,"props":1749,"children":1750},{},[1751],{"type":49,"value":1752},"# out.shape == (batch, 1 + 3 + 5 + 7)  -- sum of 2l+1\n",{"type":43,"tag":59,"props":1754,"children":1755},{},[1756,1757,1762],{"type":49,"value":1244},{"type":43,"tag":63,"props":1758,"children":1760},{"className":1759},[],[1761],{"type":49,"value":334},{"type":49,"value":975},{"type":43,"tag":541,"props":1764,"children":1766},{"id":1765},"rotation-and-inversion",[1767],{"type":49,"value":1768},"Rotation and Inversion",{"type":43,"tag":394,"props":1770,"children":1772},{"className":396,"code":1771,"language":398,"meta":399,"style":399},"# Rotation (SO3 or O3 irreps)\nrot = cuet.Rotation(\n    cue.Irreps(\"SO3\", \"4x0 + 2x1 + 1x2\"),\n    layout=cue.ir_mul,\n    device=\"cuda\",\n)\n\n# Euler angles (YXY convention)\ngamma = torch.tensor([0.1], device=\"cuda\")\nbeta = torch.tensor([0.2], device=\"cuda\")\nalpha = torch.tensor([0.3], device=\"cuda\")\nout = rot(gamma, beta, alpha, x)\n\n# Helper: encode angle for rotation\nencoded = cuet.encode_rotation_angle(angle, ell=3)  # cos\u002Fsin encoding\n\n# Helper: 3D vector to Euler angles\nbeta, alpha = cuet.vector_to_euler_angles(vector)\n\n# Inversion (O3 irreps only)\ninv = cuet.Inversion(\n    cue.Irreps(\"O3\", \"4x0e + 2x1o\"),\n    layout=cue.ir_mul,\n    device=\"cuda\",\n)\nout = inv(x)\n",[1773],{"type":43,"tag":63,"props":1774,"children":1775},{"__ignoreMap":399},[1776,1784,1792,1800,1807,1814,1821,1828,1836,1844,1852,1860,1868,1875,1883,1891,1898,1906,1914,1921,1929,1937,1945,1952,1959,1966],{"type":43,"tag":405,"props":1777,"children":1778},{"class":407,"line":408},[1779],{"type":43,"tag":405,"props":1780,"children":1781},{},[1782],{"type":49,"value":1783},"# Rotation (SO3 or O3 irreps)\n",{"type":43,"tag":405,"props":1785,"children":1786},{"class":407,"line":417},[1787],{"type":43,"tag":405,"props":1788,"children":1789},{},[1790],{"type":49,"value":1791},"rot = cuet.Rotation(\n",{"type":43,"tag":405,"props":1793,"children":1794},{"class":407,"line":426},[1795],{"type":43,"tag":405,"props":1796,"children":1797},{},[1798],{"type":49,"value":1799},"    cue.Irreps(\"SO3\", \"4x0 + 2x1 + 1x2\"),\n",{"type":43,"tag":405,"props":1801,"children":1802},{"class":407,"line":435},[1803],{"type":43,"tag":405,"props":1804,"children":1805},{},[1806],{"type":49,"value":1603},{"type":43,"tag":405,"props":1808,"children":1809},{"class":407,"line":445},[1810],{"type":43,"tag":405,"props":1811,"children":1812},{},[1813],{"type":49,"value":1055},{"type":43,"tag":405,"props":1815,"children":1816},{"class":407,"line":454},[1817],{"type":43,"tag":405,"props":1818,"children":1819},{},[1820],{"type":49,"value":592},{"type":43,"tag":405,"props":1822,"children":1823},{"class":407,"line":463},[1824],{"type":43,"tag":405,"props":1825,"children":1826},{"emptyLinePlaceholder":439},[1827],{"type":49,"value":442},{"type":43,"tag":405,"props":1829,"children":1830},{"class":407,"line":472},[1831],{"type":43,"tag":405,"props":1832,"children":1833},{},[1834],{"type":49,"value":1835},"# Euler angles (YXY convention)\n",{"type":43,"tag":405,"props":1837,"children":1838},{"class":407,"line":480},[1839],{"type":43,"tag":405,"props":1840,"children":1841},{},[1842],{"type":49,"value":1843},"gamma = torch.tensor([0.1], device=\"cuda\")\n",{"type":43,"tag":405,"props":1845,"children":1846},{"class":407,"line":489},[1847],{"type":43,"tag":405,"props":1848,"children":1849},{},[1850],{"type":49,"value":1851},"beta = torch.tensor([0.2], device=\"cuda\")\n",{"type":43,"tag":405,"props":1853,"children":1854},{"class":407,"line":498},[1855],{"type":43,"tag":405,"props":1856,"children":1857},{},[1858],{"type":49,"value":1859},"alpha = torch.tensor([0.3], device=\"cuda\")\n",{"type":43,"tag":405,"props":1861,"children":1862},{"class":407,"line":506},[1863],{"type":43,"tag":405,"props":1864,"children":1865},{},[1866],{"type":49,"value":1867},"out = rot(gamma, beta, alpha, x)\n",{"type":43,"tag":405,"props":1869,"children":1870},{"class":407,"line":515},[1871],{"type":43,"tag":405,"props":1872,"children":1873},{"emptyLinePlaceholder":439},[1874],{"type":49,"value":442},{"type":43,"tag":405,"props":1876,"children":1877},{"class":407,"line":524},[1878],{"type":43,"tag":405,"props":1879,"children":1880},{},[1881],{"type":49,"value":1882},"# Helper: encode angle for rotation\n",{"type":43,"tag":405,"props":1884,"children":1885},{"class":407,"line":533},[1886],{"type":43,"tag":405,"props":1887,"children":1888},{},[1889],{"type":49,"value":1890},"encoded = cuet.encode_rotation_angle(angle, ell=3)  # cos\u002Fsin encoding\n",{"type":43,"tag":405,"props":1892,"children":1893},{"class":407,"line":670},[1894],{"type":43,"tag":405,"props":1895,"children":1896},{"emptyLinePlaceholder":439},[1897],{"type":49,"value":442},{"type":43,"tag":405,"props":1899,"children":1900},{"class":407,"line":678},[1901],{"type":43,"tag":405,"props":1902,"children":1903},{},[1904],{"type":49,"value":1905},"# Helper: 3D vector to Euler angles\n",{"type":43,"tag":405,"props":1907,"children":1908},{"class":407,"line":687},[1909],{"type":43,"tag":405,"props":1910,"children":1911},{},[1912],{"type":49,"value":1913},"beta, alpha = cuet.vector_to_euler_angles(vector)\n",{"type":43,"tag":405,"props":1915,"children":1916},{"class":407,"line":696},[1917],{"type":43,"tag":405,"props":1918,"children":1919},{"emptyLinePlaceholder":439},[1920],{"type":49,"value":442},{"type":43,"tag":405,"props":1922,"children":1923},{"class":407,"line":705},[1924],{"type":43,"tag":405,"props":1925,"children":1926},{},[1927],{"type":49,"value":1928},"# Inversion (O3 irreps only)\n",{"type":43,"tag":405,"props":1930,"children":1931},{"class":407,"line":713},[1932],{"type":43,"tag":405,"props":1933,"children":1934},{},[1935],{"type":49,"value":1936},"inv = cuet.Inversion(\n",{"type":43,"tag":405,"props":1938,"children":1939},{"class":407,"line":722},[1940],{"type":43,"tag":405,"props":1941,"children":1942},{},[1943],{"type":49,"value":1944},"    cue.Irreps(\"O3\", \"4x0e + 2x1o\"),\n",{"type":43,"tag":405,"props":1946,"children":1947},{"class":407,"line":731},[1948],{"type":43,"tag":405,"props":1949,"children":1950},{},[1951],{"type":49,"value":1603},{"type":43,"tag":405,"props":1953,"children":1954},{"class":407,"line":740},[1955],{"type":43,"tag":405,"props":1956,"children":1957},{},[1958],{"type":49,"value":1055},{"type":43,"tag":405,"props":1960,"children":1961},{"class":407,"line":749},[1962],{"type":43,"tag":405,"props":1963,"children":1964},{},[1965],{"type":49,"value":592},{"type":43,"tag":405,"props":1967,"children":1968},{"class":407,"line":758},[1969],{"type":43,"tag":405,"props":1970,"children":1971},{},[1972],{"type":49,"value":1973},"out = inv(x)\n",{"type":43,"tag":52,"props":1975,"children":1977},{"id":1976},"layers",[1978],{"type":49,"value":192},{"type":43,"tag":541,"props":1980,"children":1982},{"id":1981},"batchnorm",[1983],{"type":49,"value":1984},"BatchNorm",{"type":43,"tag":59,"props":1986,"children":1987},{},[1988],{"type":49,"value":1989},"Batch normalization for equivariant representations (adapted from e3nn).",{"type":43,"tag":394,"props":1991,"children":1993},{"className":396,"code":1992,"language":398,"meta":399,"style":399},"bn = cuet.layers.BatchNorm(\n    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n    layout=cue.mul_ir,\n    eps=1e-5,\n    momentum=0.1,\n    affine=True,\n)\n\nout = bn(x)  # x.shape == (batch, irreps.dim)\n",[1994],{"type":43,"tag":63,"props":1995,"children":1996},{"__ignoreMap":399},[1997,2005,2013,2020,2028,2036,2044,2051,2058],{"type":43,"tag":405,"props":1998,"children":1999},{"class":407,"line":408},[2000],{"type":43,"tag":405,"props":2001,"children":2002},{},[2003],{"type":49,"value":2004},"bn = cuet.layers.BatchNorm(\n",{"type":43,"tag":405,"props":2006,"children":2007},{"class":407,"line":417},[2008],{"type":43,"tag":405,"props":2009,"children":2010},{},[2011],{"type":49,"value":2012},"    cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n",{"type":43,"tag":405,"props":2014,"children":2015},{"class":407,"line":426},[2016],{"type":43,"tag":405,"props":2017,"children":2018},{},[2019],{"type":49,"value":1047},{"type":43,"tag":405,"props":2021,"children":2022},{"class":407,"line":435},[2023],{"type":43,"tag":405,"props":2024,"children":2025},{},[2026],{"type":49,"value":2027},"    eps=1e-5,\n",{"type":43,"tag":405,"props":2029,"children":2030},{"class":407,"line":445},[2031],{"type":43,"tag":405,"props":2032,"children":2033},{},[2034],{"type":49,"value":2035},"    momentum=0.1,\n",{"type":43,"tag":405,"props":2037,"children":2038},{"class":407,"line":454},[2039],{"type":43,"tag":405,"props":2040,"children":2041},{},[2042],{"type":49,"value":2043},"    affine=True,\n",{"type":43,"tag":405,"props":2045,"children":2046},{"class":407,"line":463},[2047],{"type":43,"tag":405,"props":2048,"children":2049},{},[2050],{"type":49,"value":592},{"type":43,"tag":405,"props":2052,"children":2053},{"class":407,"line":472},[2054],{"type":43,"tag":405,"props":2055,"children":2056},{"emptyLinePlaceholder":439},[2057],{"type":49,"value":442},{"type":43,"tag":405,"props":2059,"children":2060},{"class":407,"line":480},[2061],{"type":43,"tag":405,"props":2062,"children":2063},{},[2064],{"type":49,"value":2065},"out = bn(x)  # x.shape == (batch, irreps.dim)\n",{"type":43,"tag":541,"props":2067,"children":2069},{"id":2068},"fullyconnectedtensorproductconv",[2070],{"type":49,"value":2071},"FullyConnectedTensorProductConv",{"type":43,"tag":59,"props":2073,"children":2074},{},[2075],{"type":49,"value":2076},"Message passing layer for equivariant GNNs (DiffDock-style).",{"type":43,"tag":394,"props":2078,"children":2080},{"className":396,"code":2079,"language":398,"meta":399,"style":399},"conv = cuet.layers.FullyConnectedTensorProductConv(\n    in_irreps=cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n    sh_irreps=cue.Irreps(\"O3\", \"0e + 1o\"),\n    out_irreps=cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n    mlp_channels=[16, 32, 32],\n    mlp_activation=torch.nn.ReLU(),\n    batch_norm=True,\n    layout=cue.ir_mul,\n)\n\n# graph = ((src, dst), (num_src_nodes, num_dst_nodes))\ngraph = ((src, dst), (num_src_nodes, num_dst_nodes))\n\nout = conv(src_features, edge_sh, edge_emb, graph, reduce=\"mean\")\n# out.shape == (num_dst_nodes, out_irreps.dim)\n\n# Optional: separate scalar features for efficient first-layer GEMM\nout = conv(src_features, edge_sh, edge_emb, graph,\n           src_scalars=src_scalars, dst_scalars=dst_scalars)\n",[2081],{"type":43,"tag":63,"props":2082,"children":2083},{"__ignoreMap":399},[2084,2092,2100,2108,2116,2124,2132,2140,2147,2154,2161,2169,2177,2184,2192,2200,2207,2215,2223],{"type":43,"tag":405,"props":2085,"children":2086},{"class":407,"line":408},[2087],{"type":43,"tag":405,"props":2088,"children":2089},{},[2090],{"type":49,"value":2091},"conv = cuet.layers.FullyConnectedTensorProductConv(\n",{"type":43,"tag":405,"props":2093,"children":2094},{"class":407,"line":417},[2095],{"type":43,"tag":405,"props":2096,"children":2097},{},[2098],{"type":49,"value":2099},"    in_irreps=cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n",{"type":43,"tag":405,"props":2101,"children":2102},{"class":407,"line":426},[2103],{"type":43,"tag":405,"props":2104,"children":2105},{},[2106],{"type":49,"value":2107},"    sh_irreps=cue.Irreps(\"O3\", \"0e + 1o\"),\n",{"type":43,"tag":405,"props":2109,"children":2110},{"class":407,"line":435},[2111],{"type":43,"tag":405,"props":2112,"children":2113},{},[2114],{"type":49,"value":2115},"    out_irreps=cue.Irreps(\"O3\", \"4x0e + 4x1o\"),\n",{"type":43,"tag":405,"props":2117,"children":2118},{"class":407,"line":445},[2119],{"type":43,"tag":405,"props":2120,"children":2121},{},[2122],{"type":49,"value":2123},"    mlp_channels=[16, 32, 32],\n",{"type":43,"tag":405,"props":2125,"children":2126},{"class":407,"line":454},[2127],{"type":43,"tag":405,"props":2128,"children":2129},{},[2130],{"type":49,"value":2131},"    mlp_activation=torch.nn.ReLU(),\n",{"type":43,"tag":405,"props":2133,"children":2134},{"class":407,"line":463},[2135],{"type":43,"tag":405,"props":2136,"children":2137},{},[2138],{"type":49,"value":2139},"    batch_norm=True,\n",{"type":43,"tag":405,"props":2141,"children":2142},{"class":407,"line":472},[2143],{"type":43,"tag":405,"props":2144,"children":2145},{},[2146],{"type":49,"value":1603},{"type":43,"tag":405,"props":2148,"children":2149},{"class":407,"line":480},[2150],{"type":43,"tag":405,"props":2151,"children":2152},{},[2153],{"type":49,"value":592},{"type":43,"tag":405,"props":2155,"children":2156},{"class":407,"line":489},[2157],{"type":43,"tag":405,"props":2158,"children":2159},{"emptyLinePlaceholder":439},[2160],{"type":49,"value":442},{"type":43,"tag":405,"props":2162,"children":2163},{"class":407,"line":498},[2164],{"type":43,"tag":405,"props":2165,"children":2166},{},[2167],{"type":49,"value":2168},"# graph = ((src, dst), (num_src_nodes, num_dst_nodes))\n",{"type":43,"tag":405,"props":2170,"children":2171},{"class":407,"line":506},[2172],{"type":43,"tag":405,"props":2173,"children":2174},{},[2175],{"type":49,"value":2176},"graph = ((src, dst), (num_src_nodes, num_dst_nodes))\n",{"type":43,"tag":405,"props":2178,"children":2179},{"class":407,"line":515},[2180],{"type":43,"tag":405,"props":2181,"children":2182},{"emptyLinePlaceholder":439},[2183],{"type":49,"value":442},{"type":43,"tag":405,"props":2185,"children":2186},{"class":407,"line":524},[2187],{"type":43,"tag":405,"props":2188,"children":2189},{},[2190],{"type":49,"value":2191},"out = conv(src_features, edge_sh, edge_emb, graph, reduce=\"mean\")\n",{"type":43,"tag":405,"props":2193,"children":2194},{"class":407,"line":533},[2195],{"type":43,"tag":405,"props":2196,"children":2197},{},[2198],{"type":49,"value":2199},"# out.shape == (num_dst_nodes, out_irreps.dim)\n",{"type":43,"tag":405,"props":2201,"children":2202},{"class":407,"line":670},[2203],{"type":43,"tag":405,"props":2204,"children":2205},{"emptyLinePlaceholder":439},[2206],{"type":49,"value":442},{"type":43,"tag":405,"props":2208,"children":2209},{"class":407,"line":678},[2210],{"type":43,"tag":405,"props":2211,"children":2212},{},[2213],{"type":49,"value":2214},"# Optional: separate scalar features for efficient first-layer GEMM\n",{"type":43,"tag":405,"props":2216,"children":2217},{"class":407,"line":687},[2218],{"type":43,"tag":405,"props":2219,"children":2220},{},[2221],{"type":49,"value":2222},"out = conv(src_features, edge_sh, edge_emb, graph,\n",{"type":43,"tag":405,"props":2224,"children":2225},{"class":407,"line":696},[2226],{"type":43,"tag":405,"props":2227,"children":2228},{},[2229],{"type":49,"value":2230},"           src_scalars=src_scalars, dst_scalars=dst_scalars)\n",{"type":43,"tag":52,"props":2232,"children":2234},{"id":2233},"triangle-operations-alphafold2-style",[2235],{"type":49,"value":2236},"Triangle operations (AlphaFold2-style)",{"type":43,"tag":59,"props":2238,"children":2239},{},[2240,2242,2248],{"type":49,"value":2241},"Require ",{"type":43,"tag":63,"props":2243,"children":2245},{"className":2244},[],[2246],{"type":49,"value":2247},"cuequivariance_ops_torch",{"type":49,"value":975},{"type":43,"tag":394,"props":2250,"children":2252},{"className":396,"code":2251,"language":398,"meta":399,"style":399},"# Triangle attention with pair bias\nkv_lengths = cuet.mask_to_kv_lengths(prefix_mask)\nout = cuet.triangle_attention(q, k, v, bias, scale=scale, kv_lengths=kv_lengths)\n# Arbitrary dense masks are supported through the fallback path:\nout = cuet.triangle_attention(q, k, v, bias, mask=holey_mask, scale=scale)\n# q, k, v: (B, N, H, Q\u002FK, D), bias: (B, 1, H, Q, K)\n\n# Triangle multiplicative update\nout = cuet.triangle_multiplicative_update(\n    x,         # (B, I, J, C)\n    mask=mask, # (B, I, J)\n    precision=cuet.TriMulPrecision.DEFAULT,\n)\n\n# Attention with pair bias (strict OpenFold3\u002FBoltz contract)\nout, _ = cuet.attention_pair_bias(\n    single_repr, pair_repr, mask, num_heads,\n    w_ln_a, b_ln_a,\n    w_proj_q, b_proj_q, w_proj_k, w_proj_v,\n    w_proj_g, w_proj_o, w_proj_z,\n    w_ln_z=w_ln_z, b_ln_z=b_ln_z,\n)\n# Optional generalized (Proteina\u002FComplexa) params, all keyword-only and None by\n# default: b_proj_k, b_proj_v, w_ln_q, b_ln_q, w_ln_k, b_ln_k, b_proj_g.\n",[2253],{"type":43,"tag":63,"props":2254,"children":2255},{"__ignoreMap":399},[2256,2264,2272,2280,2288,2296,2304,2311,2319,2327,2335,2343,2351,2358,2365,2373,2381,2389,2397,2405,2413,2421,2428,2436],{"type":43,"tag":405,"props":2257,"children":2258},{"class":407,"line":408},[2259],{"type":43,"tag":405,"props":2260,"children":2261},{},[2262],{"type":49,"value":2263},"# Triangle attention with pair bias\n",{"type":43,"tag":405,"props":2265,"children":2266},{"class":407,"line":417},[2267],{"type":43,"tag":405,"props":2268,"children":2269},{},[2270],{"type":49,"value":2271},"kv_lengths = cuet.mask_to_kv_lengths(prefix_mask)\n",{"type":43,"tag":405,"props":2273,"children":2274},{"class":407,"line":426},[2275],{"type":43,"tag":405,"props":2276,"children":2277},{},[2278],{"type":49,"value":2279},"out = cuet.triangle_attention(q, k, v, bias, scale=scale, kv_lengths=kv_lengths)\n",{"type":43,"tag":405,"props":2281,"children":2282},{"class":407,"line":435},[2283],{"type":43,"tag":405,"props":2284,"children":2285},{},[2286],{"type":49,"value":2287},"# Arbitrary dense masks are supported through the fallback path:\n",{"type":43,"tag":405,"props":2289,"children":2290},{"class":407,"line":445},[2291],{"type":43,"tag":405,"props":2292,"children":2293},{},[2294],{"type":49,"value":2295},"out = cuet.triangle_attention(q, k, v, bias, mask=holey_mask, scale=scale)\n",{"type":43,"tag":405,"props":2297,"children":2298},{"class":407,"line":454},[2299],{"type":43,"tag":405,"props":2300,"children":2301},{},[2302],{"type":49,"value":2303},"# q, k, v: (B, N, H, Q\u002FK, D), bias: (B, 1, H, Q, K)\n",{"type":43,"tag":405,"props":2305,"children":2306},{"class":407,"line":463},[2307],{"type":43,"tag":405,"props":2308,"children":2309},{"emptyLinePlaceholder":439},[2310],{"type":49,"value":442},{"type":43,"tag":405,"props":2312,"children":2313},{"class":407,"line":472},[2314],{"type":43,"tag":405,"props":2315,"children":2316},{},[2317],{"type":49,"value":2318},"# Triangle multiplicative update\n",{"type":43,"tag":405,"props":2320,"children":2321},{"class":407,"line":480},[2322],{"type":43,"tag":405,"props":2323,"children":2324},{},[2325],{"type":49,"value":2326},"out = cuet.triangle_multiplicative_update(\n",{"type":43,"tag":405,"props":2328,"children":2329},{"class":407,"line":489},[2330],{"type":43,"tag":405,"props":2331,"children":2332},{},[2333],{"type":49,"value":2334},"    x,         # (B, I, J, C)\n",{"type":43,"tag":405,"props":2336,"children":2337},{"class":407,"line":498},[2338],{"type":43,"tag":405,"props":2339,"children":2340},{},[2341],{"type":49,"value":2342},"    mask=mask, # (B, I, J)\n",{"type":43,"tag":405,"props":2344,"children":2345},{"class":407,"line":506},[2346],{"type":43,"tag":405,"props":2347,"children":2348},{},[2349],{"type":49,"value":2350},"    precision=cuet.TriMulPrecision.DEFAULT,\n",{"type":43,"tag":405,"props":2352,"children":2353},{"class":407,"line":515},[2354],{"type":43,"tag":405,"props":2355,"children":2356},{},[2357],{"type":49,"value":592},{"type":43,"tag":405,"props":2359,"children":2360},{"class":407,"line":524},[2361],{"type":43,"tag":405,"props":2362,"children":2363},{"emptyLinePlaceholder":439},[2364],{"type":49,"value":442},{"type":43,"tag":405,"props":2366,"children":2367},{"class":407,"line":533},[2368],{"type":43,"tag":405,"props":2369,"children":2370},{},[2371],{"type":49,"value":2372},"# Attention with pair bias (strict OpenFold3\u002FBoltz contract)\n",{"type":43,"tag":405,"props":2374,"children":2375},{"class":407,"line":670},[2376],{"type":43,"tag":405,"props":2377,"children":2378},{},[2379],{"type":49,"value":2380},"out, _ = cuet.attention_pair_bias(\n",{"type":43,"tag":405,"props":2382,"children":2383},{"class":407,"line":678},[2384],{"type":43,"tag":405,"props":2385,"children":2386},{},[2387],{"type":49,"value":2388},"    single_repr, pair_repr, mask, num_heads,\n",{"type":43,"tag":405,"props":2390,"children":2391},{"class":407,"line":687},[2392],{"type":43,"tag":405,"props":2393,"children":2394},{},[2395],{"type":49,"value":2396},"    w_ln_a, b_ln_a,\n",{"type":43,"tag":405,"props":2398,"children":2399},{"class":407,"line":696},[2400],{"type":43,"tag":405,"props":2401,"children":2402},{},[2403],{"type":49,"value":2404},"    w_proj_q, b_proj_q, w_proj_k, w_proj_v,\n",{"type":43,"tag":405,"props":2406,"children":2407},{"class":407,"line":705},[2408],{"type":43,"tag":405,"props":2409,"children":2410},{},[2411],{"type":49,"value":2412},"    w_proj_g, w_proj_o, w_proj_z,\n",{"type":43,"tag":405,"props":2414,"children":2415},{"class":407,"line":713},[2416],{"type":43,"tag":405,"props":2417,"children":2418},{},[2419],{"type":49,"value":2420},"    w_ln_z=w_ln_z, b_ln_z=b_ln_z,\n",{"type":43,"tag":405,"props":2422,"children":2423},{"class":407,"line":722},[2424],{"type":43,"tag":405,"props":2425,"children":2426},{},[2427],{"type":49,"value":592},{"type":43,"tag":405,"props":2429,"children":2430},{"class":407,"line":731},[2431],{"type":43,"tag":405,"props":2432,"children":2433},{},[2434],{"type":49,"value":2435},"# Optional generalized (Proteina\u002FComplexa) params, all keyword-only and None by\n",{"type":43,"tag":405,"props":2437,"children":2438},{"class":407,"line":740},[2439],{"type":43,"tag":405,"props":2440,"children":2441},{},[2442],{"type":49,"value":2443},"# default: b_proj_k, b_proj_v, w_ln_q, b_ln_q, w_ln_k, b_ln_k, b_proj_g.\n",{"type":43,"tag":52,"props":2445,"children":2447},{"id":2446},"onnx-and-tensorrt-export",[2448],{"type":49,"value":2449},"ONNX and TensorRT export",{"type":43,"tag":394,"props":2451,"children":2453},{"className":396,"code":2452,"language":398,"meta":399,"style":399},"# ONNX export\ntable = cuet.onnx_custom_translation_table()\nonnx_program = torch.onnx.export(model, inputs, custom_translation_table=table)\n\n# TensorRT plugin registration\ncuet.register_tensorrt_plugins()\n",[2454],{"type":43,"tag":63,"props":2455,"children":2456},{"__ignoreMap":399},[2457,2465,2473,2481,2488,2496],{"type":43,"tag":405,"props":2458,"children":2459},{"class":407,"line":408},[2460],{"type":43,"tag":405,"props":2461,"children":2462},{},[2463],{"type":49,"value":2464},"# ONNX export\n",{"type":43,"tag":405,"props":2466,"children":2467},{"class":407,"line":417},[2468],{"type":43,"tag":405,"props":2469,"children":2470},{},[2471],{"type":49,"value":2472},"table = cuet.onnx_custom_translation_table()\n",{"type":43,"tag":405,"props":2474,"children":2475},{"class":407,"line":426},[2476],{"type":43,"tag":405,"props":2477,"children":2478},{},[2479],{"type":49,"value":2480},"onnx_program = torch.onnx.export(model, inputs, custom_translation_table=table)\n",{"type":43,"tag":405,"props":2482,"children":2483},{"class":407,"line":435},[2484],{"type":43,"tag":405,"props":2485,"children":2486},{"emptyLinePlaceholder":439},[2487],{"type":49,"value":442},{"type":43,"tag":405,"props":2489,"children":2490},{"class":407,"line":445},[2491],{"type":43,"tag":405,"props":2492,"children":2493},{},[2494],{"type":49,"value":2495},"# TensorRT plugin registration\n",{"type":43,"tag":405,"props":2497,"children":2498},{"class":407,"line":454},[2499],{"type":43,"tag":405,"props":2500,"children":2501},{},[2502],{"type":49,"value":2503},"cuet.register_tensorrt_plugins()\n",{"type":43,"tag":52,"props":2505,"children":2507},{"id":2506},"complete-gnn-example",[2508],{"type":49,"value":2509},"Complete GNN example",{"type":43,"tag":394,"props":2511,"children":2513},{"className":396,"code":2512,"language":398,"meta":399,"style":399},"import torch\nimport cuequivariance as cue\nimport cuequivariance_torch as cuet\n\nclass SimpleGNN(torch.nn.Module):\n    def __init__(self, irreps_in, irreps_sh, irreps_out):\n        super().__init__()\n        self.tp = cuet.ChannelWiseTensorProduct(\n            irreps_in, irreps_sh, layout=cue.mul_ir,\n            shared_weights=False, device=\"cuda\",\n        )\n        self.linear = cuet.Linear(\n            self.tp.irreps_out, irreps_out,\n            layout=cue.mul_ir, internal_weights=True, device=\"cuda\",\n        )\n        self.sh = cuet.SphericalHarmonics(\n            ls=[ir.l for _, ir in irreps_sh], normalize=True, device=\"cuda\",\n        )\n\n    def forward(self, node_feats, edge_vec, edge_index, num_nodes):\n        src, dst = edge_index\n        edge_sh = self.sh(edge_vec)\n        w = torch.randn(1, self.tp.weight_numel, device=node_feats.device)\n\n        # Message: tensor product on edges with scatter to destination nodes\n        messages = self.tp(\n            node_feats, edge_sh, weight=w,\n            indices_1=src, indices_2=None,\n            indices_out=dst, size_out=num_nodes,\n        )\n        return self.linear(messages)\n",[2514],{"type":43,"tag":63,"props":2515,"children":2516},{"__ignoreMap":399},[2517,2524,2531,2538,2545,2553,2561,2569,2577,2585,2593,2601,2609,2617,2625,2632,2640,2648,2655,2662,2670,2678,2686,2694,2701,2709,2717,2725,2733,2741,2748],{"type":43,"tag":405,"props":2518,"children":2519},{"class":407,"line":408},[2520],{"type":43,"tag":405,"props":2521,"children":2522},{},[2523],{"type":49,"value":414},{"type":43,"tag":405,"props":2525,"children":2526},{"class":407,"line":417},[2527],{"type":43,"tag":405,"props":2528,"children":2529},{},[2530],{"type":49,"value":423},{"type":43,"tag":405,"props":2532,"children":2533},{"class":407,"line":426},[2534],{"type":43,"tag":405,"props":2535,"children":2536},{},[2537],{"type":49,"value":432},{"type":43,"tag":405,"props":2539,"children":2540},{"class":407,"line":435},[2541],{"type":43,"tag":405,"props":2542,"children":2543},{"emptyLinePlaceholder":439},[2544],{"type":49,"value":442},{"type":43,"tag":405,"props":2546,"children":2547},{"class":407,"line":445},[2548],{"type":43,"tag":405,"props":2549,"children":2550},{},[2551],{"type":49,"value":2552},"class SimpleGNN(torch.nn.Module):\n",{"type":43,"tag":405,"props":2554,"children":2555},{"class":407,"line":454},[2556],{"type":43,"tag":405,"props":2557,"children":2558},{},[2559],{"type":49,"value":2560},"    def __init__(self, irreps_in, irreps_sh, irreps_out):\n",{"type":43,"tag":405,"props":2562,"children":2563},{"class":407,"line":463},[2564],{"type":43,"tag":405,"props":2565,"children":2566},{},[2567],{"type":49,"value":2568},"        super().__init__()\n",{"type":43,"tag":405,"props":2570,"children":2571},{"class":407,"line":472},[2572],{"type":43,"tag":405,"props":2573,"children":2574},{},[2575],{"type":49,"value":2576},"        self.tp = cuet.ChannelWiseTensorProduct(\n",{"type":43,"tag":405,"props":2578,"children":2579},{"class":407,"line":480},[2580],{"type":43,"tag":405,"props":2581,"children":2582},{},[2583],{"type":49,"value":2584},"            irreps_in, irreps_sh, layout=cue.mul_ir,\n",{"type":43,"tag":405,"props":2586,"children":2587},{"class":407,"line":489},[2588],{"type":43,"tag":405,"props":2589,"children":2590},{},[2591],{"type":49,"value":2592},"            shared_weights=False, device=\"cuda\",\n",{"type":43,"tag":405,"props":2594,"children":2595},{"class":407,"line":498},[2596],{"type":43,"tag":405,"props":2597,"children":2598},{},[2599],{"type":49,"value":2600},"        )\n",{"type":43,"tag":405,"props":2602,"children":2603},{"class":407,"line":506},[2604],{"type":43,"tag":405,"props":2605,"children":2606},{},[2607],{"type":49,"value":2608},"        self.linear = cuet.Linear(\n",{"type":43,"tag":405,"props":2610,"children":2611},{"class":407,"line":515},[2612],{"type":43,"tag":405,"props":2613,"children":2614},{},[2615],{"type":49,"value":2616},"            self.tp.irreps_out, irreps_out,\n",{"type":43,"tag":405,"props":2618,"children":2619},{"class":407,"line":524},[2620],{"type":43,"tag":405,"props":2621,"children":2622},{},[2623],{"type":49,"value":2624},"            layout=cue.mul_ir, internal_weights=True, device=\"cuda\",\n",{"type":43,"tag":405,"props":2626,"children":2627},{"class":407,"line":533},[2628],{"type":43,"tag":405,"props":2629,"children":2630},{},[2631],{"type":49,"value":2600},{"type":43,"tag":405,"props":2633,"children":2634},{"class":407,"line":670},[2635],{"type":43,"tag":405,"props":2636,"children":2637},{},[2638],{"type":49,"value":2639},"        self.sh = cuet.SphericalHarmonics(\n",{"type":43,"tag":405,"props":2641,"children":2642},{"class":407,"line":678},[2643],{"type":43,"tag":405,"props":2644,"children":2645},{},[2646],{"type":49,"value":2647},"            ls=[ir.l for _, ir in irreps_sh], normalize=True, device=\"cuda\",\n",{"type":43,"tag":405,"props":2649,"children":2650},{"class":407,"line":687},[2651],{"type":43,"tag":405,"props":2652,"children":2653},{},[2654],{"type":49,"value":2600},{"type":43,"tag":405,"props":2656,"children":2657},{"class":407,"line":696},[2658],{"type":43,"tag":405,"props":2659,"children":2660},{"emptyLinePlaceholder":439},[2661],{"type":49,"value":442},{"type":43,"tag":405,"props":2663,"children":2664},{"class":407,"line":705},[2665],{"type":43,"tag":405,"props":2666,"children":2667},{},[2668],{"type":49,"value":2669},"    def forward(self, node_feats, edge_vec, edge_index, num_nodes):\n",{"type":43,"tag":405,"props":2671,"children":2672},{"class":407,"line":713},[2673],{"type":43,"tag":405,"props":2674,"children":2675},{},[2676],{"type":49,"value":2677},"        src, dst = edge_index\n",{"type":43,"tag":405,"props":2679,"children":2680},{"class":407,"line":722},[2681],{"type":43,"tag":405,"props":2682,"children":2683},{},[2684],{"type":49,"value":2685},"        edge_sh = self.sh(edge_vec)\n",{"type":43,"tag":405,"props":2687,"children":2688},{"class":407,"line":731},[2689],{"type":43,"tag":405,"props":2690,"children":2691},{},[2692],{"type":49,"value":2693},"        w = torch.randn(1, self.tp.weight_numel, device=node_feats.device)\n",{"type":43,"tag":405,"props":2695,"children":2696},{"class":407,"line":740},[2697],{"type":43,"tag":405,"props":2698,"children":2699},{"emptyLinePlaceholder":439},[2700],{"type":49,"value":442},{"type":43,"tag":405,"props":2702,"children":2703},{"class":407,"line":749},[2704],{"type":43,"tag":405,"props":2705,"children":2706},{},[2707],{"type":49,"value":2708},"        # Message: tensor product on edges with scatter to destination nodes\n",{"type":43,"tag":405,"props":2710,"children":2711},{"class":407,"line":758},[2712],{"type":43,"tag":405,"props":2713,"children":2714},{},[2715],{"type":49,"value":2716},"        messages = self.tp(\n",{"type":43,"tag":405,"props":2718,"children":2719},{"class":407,"line":767},[2720],{"type":43,"tag":405,"props":2721,"children":2722},{},[2723],{"type":49,"value":2724},"            node_feats, edge_sh, weight=w,\n",{"type":43,"tag":405,"props":2726,"children":2727},{"class":407,"line":776},[2728],{"type":43,"tag":405,"props":2729,"children":2730},{},[2731],{"type":49,"value":2732},"            indices_1=src, indices_2=None,\n",{"type":43,"tag":405,"props":2734,"children":2735},{"class":407,"line":1224},[2736],{"type":43,"tag":405,"props":2737,"children":2738},{},[2739],{"type":49,"value":2740},"            indices_out=dst, size_out=num_nodes,\n",{"type":43,"tag":405,"props":2742,"children":2743},{"class":407,"line":1233},[2744],{"type":43,"tag":405,"props":2745,"children":2746},{},[2747],{"type":49,"value":2600},{"type":43,"tag":405,"props":2749,"children":2751},{"class":407,"line":2750},31,[2752],{"type":43,"tag":405,"props":2753,"children":2754},{},[2755],{"type":49,"value":2756},"        return self.linear(messages)\n",{"type":43,"tag":52,"props":2758,"children":2760},{"id":2759},"key-file-locations",[2761],{"type":49,"value":2762},"Key file locations",{"type":43,"tag":269,"props":2764,"children":2765},{},[2766,2782],{"type":43,"tag":273,"props":2767,"children":2768},{},[2769],{"type":43,"tag":277,"props":2770,"children":2771},{},[2772,2777],{"type":43,"tag":281,"props":2773,"children":2774},{},[2775],{"type":49,"value":2776},"Component",{"type":43,"tag":281,"props":2778,"children":2779},{},[2780],{"type":49,"value":2781},"Path",{"type":43,"tag":297,"props":2783,"children":2784},{},[2785,2805,2828,2850,2872,2894,2914,2934,2954,2974,2994,3021,3041,3061,3078],{"type":43,"tag":277,"props":2786,"children":2787},{},[2788,2796],{"type":43,"tag":304,"props":2789,"children":2790},{},[2791],{"type":43,"tag":63,"props":2792,"children":2794},{"className":2793},[],[2795],{"type":49,"value":864},{"type":43,"tag":304,"props":2797,"children":2798},{},[2799],{"type":43,"tag":63,"props":2800,"children":2802},{"className":2801},[],[2803],{"type":49,"value":2804},"cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial.py",{"type":43,"tag":277,"props":2806,"children":2807},{},[2808,2819],{"type":43,"tag":304,"props":2809,"children":2810},{},[2811,2817],{"type":43,"tag":63,"props":2812,"children":2814},{"className":2813},[],[2815],{"type":49,"value":2816},"uniform_1d",{"type":49,"value":2818}," backend",{"type":43,"tag":304,"props":2820,"children":2821},{},[2822],{"type":43,"tag":63,"props":2823,"children":2825},{"className":2824},[],[2826],{"type":49,"value":2827},"cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_uniform_1d.py",{"type":43,"tag":277,"props":2829,"children":2830},{},[2831,2841],{"type":43,"tag":304,"props":2832,"children":2833},{},[2834,2840],{"type":43,"tag":63,"props":2835,"children":2837},{"className":2836},[],[2838],{"type":49,"value":2839},"naive",{"type":49,"value":2818},{"type":43,"tag":304,"props":2842,"children":2843},{},[2844],{"type":43,"tag":63,"props":2845,"children":2847},{"className":2846},[],[2848],{"type":49,"value":2849},"cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_naive.py",{"type":43,"tag":277,"props":2851,"children":2852},{},[2853,2863],{"type":43,"tag":304,"props":2854,"children":2855},{},[2856,2862],{"type":43,"tag":63,"props":2857,"children":2859},{"className":2858},[],[2860],{"type":49,"value":2861},"fused_tp",{"type":49,"value":2818},{"type":43,"tag":304,"props":2864,"children":2865},{},[2866],{"type":43,"tag":63,"props":2867,"children":2869},{"className":2868},[],[2870],{"type":49,"value":2871},"cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_fused_tp.py",{"type":43,"tag":277,"props":2873,"children":2874},{},[2875,2885],{"type":43,"tag":304,"props":2876,"children":2877},{},[2878,2884],{"type":43,"tag":63,"props":2879,"children":2881},{"className":2880},[],[2882],{"type":49,"value":2883},"indexed_linear",{"type":49,"value":2818},{"type":43,"tag":304,"props":2886,"children":2887},{},[2888],{"type":43,"tag":63,"props":2889,"children":2891},{"className":2890},[],[2892],{"type":49,"value":2893},"cuequivariance_torch\u002Fprimitives\u002Fsegmented_polynomial_indexed_linear.py",{"type":43,"tag":277,"props":2895,"children":2896},{},[2897,2905],{"type":43,"tag":304,"props":2898,"children":2899},{},[2900],{"type":43,"tag":63,"props":2901,"children":2903},{"className":2902},[],[2904],{"type":49,"value":141},{"type":43,"tag":304,"props":2906,"children":2907},{},[2908],{"type":43,"tag":63,"props":2909,"children":2911},{"className":2910},[],[2912],{"type":49,"value":2913},"cuequivariance_torch\u002Foperations\u002Ftp_channel_wise.py",{"type":43,"tag":277,"props":2915,"children":2916},{},[2917,2925],{"type":43,"tag":304,"props":2918,"children":2919},{},[2920],{"type":43,"tag":63,"props":2921,"children":2923},{"className":2922},[],[2924],{"type":49,"value":149},{"type":43,"tag":304,"props":2926,"children":2927},{},[2928],{"type":43,"tag":63,"props":2929,"children":2931},{"className":2930},[],[2932],{"type":49,"value":2933},"cuequivariance_torch\u002Foperations\u002Ftp_fully_connected.py",{"type":43,"tag":277,"props":2935,"children":2936},{},[2937,2945],{"type":43,"tag":304,"props":2938,"children":2939},{},[2940],{"type":43,"tag":63,"props":2941,"children":2943},{"className":2942},[],[2944],{"type":49,"value":156},{"type":43,"tag":304,"props":2946,"children":2947},{},[2948],{"type":43,"tag":63,"props":2949,"children":2951},{"className":2950},[],[2952],{"type":49,"value":2953},"cuequivariance_torch\u002Foperations\u002Flinear.py",{"type":43,"tag":277,"props":2955,"children":2956},{},[2957,2965],{"type":43,"tag":304,"props":2958,"children":2959},{},[2960],{"type":43,"tag":63,"props":2961,"children":2963},{"className":2962},[],[2964],{"type":49,"value":163},{"type":43,"tag":304,"props":2966,"children":2967},{},[2968],{"type":43,"tag":63,"props":2969,"children":2971},{"className":2970},[],[2972],{"type":49,"value":2973},"cuequivariance_torch\u002Foperations\u002Fsymmetric_contraction.py",{"type":43,"tag":277,"props":2975,"children":2976},{},[2977,2985],{"type":43,"tag":304,"props":2978,"children":2979},{},[2980],{"type":43,"tag":63,"props":2981,"children":2983},{"className":2982},[],[2984],{"type":49,"value":170},{"type":43,"tag":304,"props":2986,"children":2987},{},[2988],{"type":43,"tag":63,"props":2989,"children":2991},{"className":2990},[],[2992],{"type":49,"value":2993},"cuequivariance_torch\u002Foperations\u002Fspherical_harmonics.py",{"type":43,"tag":277,"props":2995,"children":2996},{},[2997,3012],{"type":43,"tag":304,"props":2998,"children":2999},{},[3000,3005,3007],{"type":43,"tag":63,"props":3001,"children":3003},{"className":3002},[],[3004],{"type":49,"value":177},{"type":49,"value":3006}," \u002F ",{"type":43,"tag":63,"props":3008,"children":3010},{"className":3009},[],[3011],{"type":49,"value":184},{"type":43,"tag":304,"props":3013,"children":3014},{},[3015],{"type":43,"tag":63,"props":3016,"children":3018},{"className":3017},[],[3019],{"type":49,"value":3020},"cuequivariance_torch\u002Foperations\u002Frotation.py",{"type":43,"tag":277,"props":3022,"children":3023},{},[3024,3032],{"type":43,"tag":304,"props":3025,"children":3026},{},[3027],{"type":43,"tag":63,"props":3028,"children":3030},{"className":3029},[],[3031],{"type":49,"value":1984},{"type":43,"tag":304,"props":3033,"children":3034},{},[3035],{"type":43,"tag":63,"props":3036,"children":3038},{"className":3037},[],[3039],{"type":49,"value":3040},"cuequivariance_torch\u002Flayers\u002Fbatchnorm.py",{"type":43,"tag":277,"props":3042,"children":3043},{},[3044,3052],{"type":43,"tag":304,"props":3045,"children":3046},{},[3047],{"type":43,"tag":63,"props":3048,"children":3050},{"className":3049},[],[3051],{"type":49,"value":2071},{"type":43,"tag":304,"props":3053,"children":3054},{},[3055],{"type":43,"tag":63,"props":3056,"children":3058},{"className":3057},[],[3059],{"type":49,"value":3060},"cuequivariance_torch\u002Flayers\u002Ftp_conv_fully_connected.py",{"type":43,"tag":277,"props":3062,"children":3063},{},[3064,3069],{"type":43,"tag":304,"props":3065,"children":3066},{},[3067],{"type":49,"value":3068},"Triangle operations",{"type":43,"tag":304,"props":3070,"children":3071},{},[3072],{"type":43,"tag":63,"props":3073,"children":3075},{"className":3074},[],[3076],{"type":49,"value":3077},"cuequivariance_torch\u002Fprimitives\u002Ftriangle.py",{"type":43,"tag":277,"props":3079,"children":3080},{},[3081,3086],{"type":43,"tag":304,"props":3082,"children":3083},{},[3084],{"type":49,"value":3085},"Layout transposition",{"type":43,"tag":304,"props":3087,"children":3088},{},[3089],{"type":43,"tag":63,"props":3090,"children":3092},{"className":3091},[],[3093],{"type":49,"value":3094},"cuequivariance_torch\u002Fprimitives\u002Ftranspose.py",{"type":43,"tag":3096,"props":3097,"children":3098},"style",{},[3099],{"type":49,"value":3100},"html .light .shiki span {color: var(--shiki-light);background: var(--shiki-light-bg);font-style: var(--shiki-light-font-style);font-weight: var(--shiki-light-font-weight);text-decoration: var(--shiki-light-text-decoration);}html.light .shiki span {color: var(--shiki-light);background: var(--shiki-light-bg);font-style: var(--shiki-light-font-style);font-weight: var(--shiki-light-font-weight);text-decoration: var(--shiki-light-text-decoration);}html .default .shiki span {color: var(--shiki-default);background: var(--shiki-default-bg);font-style: var(--shiki-default-font-style);font-weight: var(--shiki-default-font-weight);text-decoration: var(--shiki-default-text-decoration);}html .shiki span {color: var(--shiki-default);background: var(--shiki-default-bg);font-style: var(--shiki-default-font-style);font-weight: var(--shiki-default-font-weight);text-decoration: var(--shiki-default-text-decoration);}html .dark .shiki span {color: var(--shiki-dark);background: var(--shiki-dark-bg);font-style: var(--shiki-dark-font-style);font-weight: var(--shiki-dark-font-weight);text-decoration: var(--shiki-dark-text-decoration);}html.dark .shiki span {color: var(--shiki-dark);background: var(--shiki-dark-bg);font-style: var(--shiki-dark-font-style);font-weight: var(--shiki-dark-font-weight);text-decoration: var(--shiki-dark-text-decoration);}",{"items":3102,"total":426},[3103,3115,3127],{"slug":84,"name":84,"fn":3104,"description":3105,"org":3106,"tags":3107,"stars":26,"repoUrl":27,"updatedAt":3114},"build equivariant tensor products with cuEquivariance","Define custom groups (Irrep subclasses), build segmented tensor products with CG coefficients, create equivariant polynomials and IrDictPolynomials, and use built-in descriptors (linear, tensor products, spherical harmonics). Use when working with cuequivariance group theory, irreps, or segmented polynomials.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3108,3109,3112,3113],{"name":13,"slug":14,"type":15},{"name":3110,"slug":3111,"type":15},"Engineering","engineering",{"name":17,"slug":18,"type":15},{"name":9,"slug":8,"type":15},"2026-07-14T05:31:46.310103",{"slug":3116,"name":3116,"fn":3117,"description":3118,"org":3119,"tags":3120,"stars":26,"repoUrl":27,"updatedAt":3126},"cuequivariance-jax","execute equivariant polynomials in JAX","Execute equivariant polynomials in JAX using segmented_polynomial (naive\u002Funiform_1d), the ir_dict workflow with IrDictPolynomial and dict[Irrep, Array], and Flax NNX layers (IrrepsLinear, SphericalHarmonics, IrrepsIndexedLinear). Use when writing JAX code with cuequivariance.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3121,3122,3123,3124],{"name":13,"slug":14,"type":15},{"name":17,"slug":18,"type":15},{"name":9,"slug":8,"type":15},{"name":3125,"slug":398,"type":15},"Python","2026-07-14T05:34:29.712307",{"slug":4,"name":4,"fn":5,"description":6,"org":3128,"tags":3129,"stars":26,"repoUrl":27,"updatedAt":28},{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3130,3131,3132,3133,3134],{"name":13,"slug":14,"type":15},{"name":17,"slug":18,"type":15},{"name":9,"slug":8,"type":15},{"name":24,"slug":25,"type":15},{"name":20,"slug":21,"type":15},{"items":3136,"total":3293},[3137,3155,3171,3182,3194,3208,3221,3235,3248,3259,3273,3282],{"slug":3138,"name":3138,"fn":3139,"description":3140,"org":3141,"tags":3142,"stars":3152,"repoUrl":3153,"updatedAt":3154},"nemoclaw-user-guide","retrieve NemoClaw documentation and configuration","Guides human users' AI agents to the NemoClaw docs MCP server and canonical Fern documentation in Markdown form. Use when users ask how to install, configure, operate, troubleshoot, secure, or learn NemoClaw with an AI coding assistant. Trigger keywords - nemoclaw docs, use nemoclaw with ai agent, nemoclaw mcp docs, nemoclaw install help, nemoclaw quickstart, nemoclaw markdown docs, llms.txt, agent skills.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3143,3146,3149],{"name":3144,"slug":3145,"type":15},"Documentation","documentation",{"name":3147,"slug":3148,"type":15},"MCP","mcp",{"name":3150,"slug":3151,"type":15},"Search","search",21777,"https:\u002F\u002Fgithub.com\u002FNVIDIA\u002FNemoClaw","2026-07-20T06:00:01.461044",{"slug":3156,"name":3156,"fn":3157,"description":3158,"org":3159,"tags":3160,"stars":3168,"repoUrl":3169,"updatedAt":3170},"mcore-build-and-dependency","manage Megatron-LM development environments","Container-based dev environment setup and dependency management for Megatron-LM. Covers acquiring and launching the CI container, uv package management, and updating uv.lock.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3161,3164,3167],{"name":3162,"slug":3163,"type":15},"Containers","containers",{"name":3165,"slug":3166,"type":15},"Deployment","deployment",{"name":3125,"slug":398,"type":15},17049,"https:\u002F\u002Fgithub.com\u002FNVIDIA\u002FMegatron-LM","2026-07-27T06:06:11.249662",{"slug":3172,"name":3172,"fn":3173,"description":3174,"org":3175,"tags":3176,"stars":3168,"repoUrl":3169,"updatedAt":3181},"mcore-bump-base-image","update NVIDIA PyTorch base images","Bump the NVIDIA PyTorch base image (`nvcr.io\u002Fnvidia\u002Fpytorch:YY.MM-py3`) used by Megatron-LM CI. Covers the two pin sites (GitHub CI in `docker\u002F.ngc_version.dev` and GitLab CI in `.gitlab\u002Fstages\u002F01.build.yml`), the post-bump CI loop (re-run functional tests, refresh golden values, mark broken tests), and the gotchas that bit PRs",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3177,3180],{"name":3178,"slug":3179,"type":15},"CI\u002FCD","ci-cd",{"name":3165,"slug":3166,"type":15},"2026-07-14T05:25:59.97109",{"slug":3183,"name":3183,"fn":3184,"description":3185,"org":3186,"tags":3187,"stars":3168,"repoUrl":3169,"updatedAt":3193},"mcore-cicd","manage CI\u002FCD pipelines for Megatron-LM","CI\u002FCD reference for Megatron-LM. Covers CI pipeline structure, PR scope labels, triggering internal GitLab CI (which force-pushes the current branch to a pull-request\u002FBRANCH ref — always dry-run and verify the destination first; never run against shared or protected branches), and CI failure investigation.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3188,3189,3190],{"name":3178,"slug":3179,"type":15},{"name":3165,"slug":3166,"type":15},{"name":3191,"slug":3192,"type":15},"GitHub","github","2026-07-27T06:06:12.278222",{"slug":3195,"name":3195,"fn":3196,"description":3197,"org":3198,"tags":3199,"stars":3168,"repoUrl":3169,"updatedAt":3207},"mcore-create-issue","investigate CI failures and create issues","Investigate a failing GitHub Actions run or job and create a GitHub issue for the failure.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3200,3203,3204],{"name":3201,"slug":3202,"type":15},"Debugging","debugging",{"name":3191,"slug":3192,"type":15},{"name":3205,"slug":3206,"type":15},"Triage","triage","2026-07-14T05:25:57.442089",{"slug":3209,"name":3209,"fn":3210,"description":3211,"org":3212,"tags":3213,"stars":3168,"repoUrl":3169,"updatedAt":3220},"mcore-linting-and-formatting","lint and format Megatron-LM code","Linting and formatting for Megatron-LM. Covers running autoformat.sh, tools (ruff, black, isort, pylint, mypy), and code style rules.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3214,3217],{"name":3215,"slug":3216,"type":15},"Best Practices","best-practices",{"name":3218,"slug":3219,"type":15},"Code Analysis","code-analysis","2026-07-14T05:25:56.18433",{"slug":3222,"name":3222,"fn":3223,"description":3224,"org":3225,"tags":3226,"stars":3168,"repoUrl":3169,"updatedAt":3234},"mcore-migrate-gpt-to-hybrid","migrate Megatron-LM models to HybridModel","Migration guide for moving Megatron Core GPTModel checkpoints, model providers, training commands, and layer mappings to HybridModel.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3227,3230,3233],{"name":3228,"slug":3229,"type":15},"Machine Learning","machine-learning",{"name":3231,"slug":3232,"type":15},"Migration","migration",{"name":9,"slug":8,"type":15},"2026-07-17T06:07:11.777011",{"slug":3236,"name":3236,"fn":3237,"description":3238,"org":3239,"tags":3240,"stars":3168,"repoUrl":3169,"updatedAt":3247},"mcore-onboard-gb200-1node-tests","onboard functional tests for GB200","Onboard 1-node GitHub MR functional tests for GB200 from existing mr-scoped 2-node tests.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3241,3244],{"name":3242,"slug":3243,"type":15},"QA","qa",{"name":3245,"slug":3246,"type":15},"Testing","testing","2026-07-14T05:25:53.673039",{"slug":3249,"name":3249,"fn":3250,"description":3251,"org":3252,"tags":3253,"stars":3168,"repoUrl":3169,"updatedAt":3258},"mcore-run-on-slurm","launch distributed training jobs on SLURM","How to launch distributed Megatron-LM training jobs on a SLURM cluster. Covers a minimal sbatch skeleton, environment-variable setup for torch.distributed.run, CUDA_DEVICE_MAX_CONNECTIONS rules across hardware and parallelism modes, container conventions, monitoring, and per-rank failure diagnosis.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3254,3255],{"name":3165,"slug":3166,"type":15},{"name":3256,"slug":3257,"type":15},"Infrastructure","infrastructure","2026-07-14T05:25:49.362534",{"slug":3260,"name":3260,"fn":3261,"description":3262,"org":3263,"tags":3264,"stars":3168,"repoUrl":3169,"updatedAt":3272},"mcore-split-pr","split pull requests to reduce review load","Split a PR into multiple PRs to reduce the number of required CODEOWNERS reviewer groups.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3265,3268,3269],{"name":3266,"slug":3267,"type":15},"Code Review","code-review",{"name":3191,"slug":3192,"type":15},{"name":3270,"slug":3271,"type":15},"Pull Requests","pull-requests","2026-07-14T05:26:01.226578",{"slug":3274,"name":3274,"fn":3275,"description":3276,"org":3277,"tags":3278,"stars":3168,"repoUrl":3169,"updatedAt":3281},"mcore-testing","run and manage Megatron-LM tests","Test system for Megatron-LM. Covers test layout, recipe YAML structure, adding and running unit and functional tests, golden values, marker filters, and CI parity.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3279,3280],{"name":3242,"slug":3243,"type":15},{"name":3245,"slug":3246,"type":15},"2026-07-14T05:25:54.928983",{"slug":3283,"name":3283,"fn":3284,"description":3285,"org":3286,"tags":3287,"stars":3168,"repoUrl":3169,"updatedAt":3292},"nightly-sync","manage nightly main-to-dev sync workflows","Domain knowledge for the nightly main-to-dev sync workflow. Covers merge strategy, CI architecture, failure investigation, and known issues.",{"slug":8,"name":9,"logoUrl":10,"githubOrg":9},[3288,3291],{"name":3289,"slug":3290,"type":15},"Automation","automation",{"name":3178,"slug":3179,"type":15},"2026-07-30T05:29:03.275638",496]