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