Cuequivariance Jax

NVIDIAOfficialCreated by Mario GeigerDistributed by NVIDIA

Skill
Tool integration
Software development
Apache-2.0
Medium risk

From the author

Use when writing or debugging cuEquivariance JAX polynomials, irrep dictionaries, Flax NNX layers, or CPU/CUDA execution and layouts.

— the description by Mario Geiger

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides instructions for integrating and debugging cuEquivariance JAX (cuex) within JAX-based projects. It covers the execution of equivariant polynomials, management of irrep dictionaries, and the implementation of Flax NNX layers. The documentation details environment setup, backend selection for CPU and CUDA, and specific data layout conversions required for handling equivariant representations.

What it does

Executes cuEquivariance descriptors using JAX.
Provides recipes for Flax NNX layer implementation.
Manages irrep dictionary layouts and data conversions.
Configures CPU and CUDA execution backends.
Supports debugging of polynomial operations and numerical mismatches.
Facilitates integration with JAX-based GNN workflows.

When to use it

When developing or debugging equivariant neural networks in JAX.
When implementing custom Flax NNX layers using cuEquivariance.
When performing tensor product operations with specific irrep layouts.
When troubleshooting numerical precision or layout issues in equivariant models.

Advantages

Enables consistent handling of equivariant representations in JAX.
Provides clear patterns for layout transformations and dictionary management.
Supports both naive CPU execution and optimized CUDA backends.
Integrates directly with Flax NNX for modular neural network design.

Usage examples

“How do I implement a basic IrrepsLinear layer using Flax NNX?”: The agent provides a code example using cue_nnx.IrrepsLinear, demonstrating how to initialize the layer with input/output irreps, handle RNGs, and perform forward passes with equivariant features.
“How can I convert between ir_mul and mul_ir layouts in my irrep dictionary?”: The agent provides functions to transpose blocks and regroup data, ensuring that round trips between flat arrays and dictionary formats maintain correct alignment.
“What is the correct way to run a segmented polynomial on a CPU?”: The agent provides a script using method='naive' within cuex.segmented_polynomial, including necessary JAX configuration for float64 and proper broadcasting of input axes.

Install

Download the skill as a ZIP:
https://autogineer.com/agent-skills/nvidia/cuequivariance-jax/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip cuequivariance-jax.zip -d .agents/skills

Documentation by Mario Geiger

cuEquivariance in JAX

Purpose

Execute cuequivariance (cue) descriptors with cuequivariance_jax (cuex). Preserve the project's framework and devices. Custom representations use core cuequivariance; PyTorch modules use cuequivariance_torch.

Inputs

Required inputs: operation, SO3/O3 irreps, layout, batch shapes, dtype, device, and weight ownership. Derive them from the request/code; state defaults. Graphs also need indices and destination count. Operand sizes/output irreps come from descriptors. Optional inputs: backend/performance preferences.

Instructions

Probe the project interpreter once and reuse working packages. Environment setup/recovery must not recursively delete directories, including temporary environments. Use a fresh unique path when isolation is needed.
Select method="naive" for CPU; Flax is needed only for NNX. High-level dictionary blocks are (..., mul, ir.dim); uniform CUDA arrays differ.
Adapt one matching recipe into a script. Build descriptors outside transformed functions, broadcast leading input axes, and enable x64 before array creation for float64 checks.
Finish installation before running the saved script with the completion recipe below. A running session/task ID, empty output, or a file write is not a completed check. Poll the same session until it exits; then read its results before reporting success.
Investigate a concrete failure using its error and the affected installed API. Rerun affected checks after a fix. Once the requested checks pass, finish; avoid repeated probes, unrelated examples, or exploration of other backends for a CPU reference.

Prerequisites and environment recovery

Probe imports independently so a missing frontend does not hide working JAX. For an NNX task, include "flax.nnx" in the same loop:

import importlib
import sys

print("python:", sys.executable, sys.version.split()[0])
available, import_errors = {}, {}
for name in ("jax", "cuequivariance", "cuequivariance_jax"):
    try:
        module = importlib.import_module(name)
        available[name] = module
        print(name, getattr(module, "__version__", "unknown"), module.__file__)
    except (ImportError, OSError) as error:
        import_errors[name] = f"{type(error).__name__}: {error}"
print("import errors:", import_errors)
if "jax" in available:
    print("devices:", available["jax"].devices())

Reuse working packages; install only missing dependencies when needed. Honor the project's lockfile or supplied scripts/requirements.txt and compatible core/frontend revisions. JAX 0.11.1 / Flax 0.12.10 require Python >=3.12: select a compatible interpreter for those pins. The standalone combination below also supports Python 3.11; it does not replace project pins. Imports alone do not establish compatibility: cuEquivariance 0.13.0 with JAX 0.12.0 fails polynomial execution with a segmented_polynomial_dce argument-count error.

For isolation, keep the unique directory returned by mktemp; reuse it for reruns. This Linux recipe uses host pip 22.3+ and bypasses ensurepip:

set -e
cueq_python=python3  # Set to the selected project interpreter's path.
cueq_jax_env=$(mktemp -d "${TMPDIR:-/tmp}/cueq-jax.XXXXXX")
"$cueq_python" -m venv --without-pip "$cueq_jax_env"
python3 -m pip --python "$cueq_jax_env" install --cert "${PIP_CERT:-/etc/ssl/certs/ca-certificates.crt}" -q \
    "jax==0.10.2" "cuequivariance==0.13.0" "cuequivariance-jax==0.13.0"
cueq_python="$cueq_jax_env/bin/python"
JAX_PLATFORMS=cpu "$cueq_python" -c 'import jax, cuequivariance_jax; print("setup complete:", jax.__version__, cuequivariance_jax.__version__, jax.devices())'

Replace the package list with -r and the requirements path when prescribed; otherwise add flax==0.12.7 to the same install command only for NNX. Keep this interpreter for execution; shell variables may need restoring in a new tool call. Without host pip, use an installed uv. The pip command uses the Linux system CA bundle; for a private CA use the environment's approved PIP_CERT bundle. Retain TLS verification. Diagnose the actual error before retrying; never reset directories or switch versions while an install is still running. If no compatible interpreter or trusted package source is available, report the blocker and unrun checks. See pip's certificates.

Complete and inspect a run

Save the preamble and relevant checks in check.py. Set cueq_python to the probe's sys.executable, or retain it from setup. Run after setup exits successfully; a fresh directory records output and status:

set -e
cueq_run=$(mktemp -d "${TMPDIR:-/tmp}/cueq-check.XXXXXX")
printf 'run directory: %s\n' "$cueq_run"
if JAX_PLATFORMS=cpu "$cueq_python" -u check.py >"$cueq_run/output.txt" 2>&1; then
    cueq_status=0
else
    cueq_status=$?
fi
printf '%s\n' "$cueq_status" >"$cueq_run/exit-code.txt"
tail -n 30 "$cueq_run/output.txt"
printf 'exit code: %s\n' "$cueq_status"
test "$cueq_status" -eq 0

If the tool yields, keep polling its session; the exit-code file is written only after Python finishes. Inspect that file and the numerical output. Empty output is not evidence that assertions ran. Deliver the script, actual command, and observed checks; label anything unrun. Keep setup output separate from the numerical results.

Limitations and backend selection

Route	Supported use
`segmented_polynomial(..., method="naive")`	Portable JAX execution, including CPU, JIT, gradients, and vmap.
`cuex.nnx.IrrepsLinear`	CPU/GPU channel mixing within matching irreps using JAX einsum.
`cuex.nnx.SphericalHarmonics`	CPU/GPU; normalizes vectors and uses the naive polynomial implementation.
`method="uniform_1d"`	CUDA; uniform segment shapes within each operand and a single mode after preparation.
`method="indexed_linear"`, `cuex.nnx.IrrepsIndexedLinear`	CUDA indexed linear route, using `cuex.Repeats`; not the ordinary CPU linear layer.

CPU routes need no GPU dependencies; moving arrays to CPU does not change the method. CUDA requires compatible JAX/GPU/kernel builds. NNX spherical harmonics normalize vectors: use positive eps if zero vectors are possible.

Examples

Use this preamble with one recipe and run with JAX_PLATFORMS=cpu. Nonconstant inputs expose layout errors; prints follow successful assertions.

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
import cuequivariance as cue
import cuequivariance_jax as cuex
from cuequivariance_jax import ir_dict

dtype = jnp.float64
print("devices:", jax.devices())

Spherical harmonics without Flax

This evaluates the raw polynomial; normalize vectors for unit-direction harmonics. O3 keys retain parity; SO3 keys do not.

desc = cue.descriptors.spherical_harmonics_ir_dict(cue.O3(1, -1), [0, 1, 2])
vectors = jnp.array([[0.3, -0.7, 1.1], [1.2, 0.4, -0.2]], dtype=dtype)
values = cuex.segmented_polynomial(
    desc.polynomial, [vectors],
    [jax.ShapeDtypeStruct(vectors.shape[:-1] + (op.size,), dtype)
     for op in desc.polynomial.outputs],
    method="naive",
)
(irreps_out,) = desc.output_irreps
harmonics = {
    ir: value.reshape(vectors.shape[:-1] + (mul, ir.dim))
    for (mul, ir), value in zip(irreps_out, values)
}
for actual, expected in zip(values, desc.polynomial(np.asarray(vectors))):
    np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12)
assert {ir.l: v.shape for ir, v in harmonics.items()} == {
    0: (2, 1, 1), 1: (2, 1, 3), 2: (2, 1, 5),
}
print({str(ir): (v.shape, v.devices()) for ir, v in harmonics.items()})

Dictionary layouts and exact round trips

flat_to_dict(..., layout="ir_mul") produces (..., mul, ir.dim) blocks. dict_to_flat has no layout argument and emits mul_ir. Return to ir_mul by transposing each block, not reshaping the whole array.

Explain the unique-irrep key precondition, even for regrouped inputs. Regroup duplicate blocks with their data using RepArray.regroup(); changing only metadata corrupts alignment. The round trip below recovers the regrouped buffer. Recovering an ungrouped order also needs its original block arrangement.

def ir_mul_to_dict(irreps, flat):
    # Merge duplicate irrep keys and reorder their data together.
    data = cuex.RepArray(cue.IrrepsAndLayout(irreps, cue.ir_mul), flat).regroup()
    return data.irreps, ir_dict.flat_to_dict(data.irreps, data.array, layout="ir_mul")

def dict_to_ir_mul_flat(irreps, blocks):
    return jnp.concatenate([
        blocks[ir].swapaxes(-1, -2).reshape(
            blocks[ir].shape[:-2] + (mul * ir.dim,)
        )
        for mul, ir in irreps
    ], axis=-1)

irreps = cue.Irreps("SO3", "3x0 + 2x1").regroup()
flat_ir = jnp.arange(2 * irreps.dim).reshape(2, irreps.dim)
irreps, blocks = ir_mul_to_dict(irreps, flat_ir)
offset = 0
for mul, ir in irreps:
    size = mul * ir.dim
    expected = np.asarray(flat_ir[..., offset:offset + size]).reshape(2, ir.dim, mul)
    assert blocks[ir].shape == (2, mul, ir.dim)
    np.testing.assert_array_equal(blocks[ir], expected.swapaxes(-1, -2))
    offset += size
flat_mul = ir_dict.dict_to_flat(irreps, blocks)
restored_blocks = ir_dict.flat_to_dict(irreps, flat_mul, layout="mul_ir")
for ir in blocks:
    np.testing.assert_array_equal(restored_blocks[ir], blocks[ir])
np.testing.assert_array_equal(dict_to_ir_mul_flat(irreps, restored_blocks), flat_ir)
print("regrouped unique irrep keys; exact ir_mul/mul_ir round trip:",
      {str(ir): v.shape for ir, v in blocks.items()})

ir_dict.assert_mul_ir_dict checks block conventions. mul_ir_dict(irreps, data) broadcasts a template to keys without inferring feature dimensions.

Dense polynomial: broadcasting, JIT, and a gradient

segmented_polynomial takes/returns lists. Channelwise products put weights first; shared weights need no batch axis. Broadcast leading input axes. filter_irreps_out selects types, not multiplicities; e.inputs/e.outputs carry dense operand layouts.

e = cue.descriptors.channelwise_tensor_product(
    cue.Irreps("SO3", "3x0 + 3x1"), cue.Irreps("SO3", "0 + 1"),
    [cue.SO3(0), cue.SO3(1)],
)
poly = e.polynomial
rng = np.random.default_rng(7)
w = jnp.asarray(rng.normal(size=poly.inputs[0].size))
x = jnp.asarray(rng.normal(size=(3, 1, poly.inputs[1].size)))
y = jnp.asarray(rng.normal(size=(1, 2, poly.inputs[2].size)))

def run(weights, a, b):
    batch = np.broadcast_shapes(weights.shape[:-1], a.shape[:-1], b.shape[:-1])
    return cuex.segmented_polynomial(
        poly, [weights, a, b],
        [jax.ShapeDtypeStruct(batch + (op.size,), jnp.result_type(weights, a, b))
         for op in poly.outputs],
        method="naive",
    )[0]

reference = e(*(np.asarray(v) for v in (w, x, y)))[0]
eager, compiled = run(w, x, y), jax.jit(run)(w, x, y)
np.testing.assert_allclose(eager, reference, rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(compiled, reference, rtol=1e-12, atol=1e-12)
assert reference.shape == (3, 2, poly.outputs[0].size)
gradient = jax.grad(lambda v: jnp.sum(run(v, x, y) ** 2))(w)
index = int(jnp.argmax(jnp.abs(gradient)))
assert abs(gradient[index]) > 1e-8
epsilon = 1e-5
loss_plus = np.sum(e(np.asarray(w.at[index].add(epsilon)), np.asarray(x), np.asarray(y))[0] ** 2)
loss_minus = np.sum(e(np.asarray(w.at[index].add(-epsilon)), np.asarray(x), np.asarray(y))[0] ** 2)
finite_difference = (loss_plus - loss_minus) / (2 * epsilon)
np.testing.assert_allclose(gradient[index], finite_difference, rtol=1e-6, atol=1e-7)
print("shape:", eager.shape, "eager/JIT max errors:",
      float(jnp.max(abs(eager - reference))), float(jnp.max(abs(compiled - reference))))
print("weight gradient / finite difference:", float(gradient[index]), finite_difference)

For gather/scatter, supply one indices entry per operand, inputs then outputs; None means unindexed. Output indices need a destination shape; repeated destinations accumulate. Verify against explicit edgewise summation.

Split irrep polynomials on CPU

channelwise_tensor_product_ir_dict returns an IrDictPolynomial. Follow input_irreps order: weights, features, second input; each (mul, ir) is one operand. Follow output_irreps for outputs. These descriptors use ir_mul: transpose blocks before flattening; reverse on outputs. There is no ir_dict.equivariant_polynomial API.

def split_naive(desc, groups):
    assert len(groups) == len(desc.input_irreps)
    operands = [
        group[ir].swapaxes(-1, -2).reshape(group[ir].shape[:-2] + (mul * ir.dim,))
        for irreps, group in zip(desc.input_irreps, groups)
        for mul, ir in irreps
    ]
    batch = np.broadcast_shapes(*(v.shape[:-1] for v in operands))
    outputs = cuex.segmented_polynomial(
        desc.polynomial, operands,
        [jax.ShapeDtypeStruct(batch + (op.size,), jnp.result_type(*operands))
         for op in desc.polynomial.outputs],
        method="naive",
    )
    values = iter(outputs)
    return tuple({
        ir: next(values).reshape(batch + (ir.dim, mul)).swapaxes(-1, -2)
        for mul, ir in irreps
    } for irreps in desc.output_irreps)

desc = cue.descriptors.channelwise_tensor_product_ir_dict(
    cue.Irreps("SO3", "3x0 + 3x1"), cue.Irreps("SO3", "0 + 1"),
    [cue.SO3(0), cue.SO3(1)],
)
rng = np.random.default_rng(8)
groups = [
    {ir: jnp.asarray(rng.normal(size=batch + (mul, ir.dim))) for mul, ir in irreps}
    for irreps, batch in zip(desc.input_irreps, [(), (3, 1), (1, 2)])
]
(result,) = split_naive(desc, groups)
numpy_inputs = [
    np.asarray(group[ir]).swapaxes(-1, -2).reshape(group[ir].shape[:-2] + (mul * ir.dim,))
    for irreps, group in zip(desc.input_irreps, groups) for mul, ir in irreps
]
expected = desc.polynomial(*numpy_inputs)
for (mul, ir), flat in zip(desc.output_irreps[0], expected):
    assert result[ir].shape == (3, 2, mul, ir.dim)
    reference = flat.reshape(3, 2, ir.dim, mul).swapaxes(-1, -2)
    np.testing.assert_allclose(result[ir], reference, rtol=1e-12, atol=1e-12)
    print("output:", str(ir), result[ir].shape, "max error:",
          float(jnp.max(abs(result[ir] - reference))))

Flax NNX layers

IrrepsLinear requires regrouped irreps and nnx.Rngs. It mixes matching-irrep channels with scale / sqrt(mul_in) normalization; unmatched outputs are zero. Read nnx.Param weights with [...]. Row vectors rotate as x @ D.T.

from flax import nnx
from cuequivariance_jax import nnx as cue_nnx

irreps_in = cue.Irreps("SO3", "3x0 + 2x1").regroup()
irreps_out = cue.Irreps("SO3", "2x0 + 4x1 + 2").regroup()
linear = cue_nnx.IrrepsLinear(irreps_in, irreps_out, scale=0.7, dtype=dtype, rngs=nnx.Rngs(9))
rng = np.random.default_rng(10)
features = {ir: jnp.asarray(rng.normal(size=(2, mul, ir.dim))) for mul, ir in irreps_in}
output = linear(features)
for mul, ir in irreps_out:
    assert output[ir].shape == (2, mul, ir.dim)
np.testing.assert_array_equal(output[cue.SO3(2)], 0)
for mul, ir in irreps_in:
    factor = linear.scale / np.sqrt(mul)
    expected = factor * np.einsum(
        "uv,...ui->...vi", np.asarray(linear.w[str(ir)][...]), features[ir])
    np.testing.assert_allclose(output[ir], expected, rtol=1e-12, atol=1e-12)
    print("irrep:", str(ir), "scale/sqrt(mul_in):", factor,
          "reference max error:", float(jnp.max(abs(output[ir] - expected))))
rotation = np.array([0.3, -0.2, 0.1])
rotated_output = linear({ir: v @ ir.exp_map(rotation, []).T for ir, v in features.items()})
for ir, value in output.items():
    expected = value @ ir.exp_map(rotation, []).T
    np.testing.assert_allclose(rotated_output[ir], expected, rtol=1e-12, atol=1e-12)
    print("output:", str(ir), value.shape, "rotation max error:",
          float(jnp.max(abs(rotated_output[ir] - expected))))
print("unmatched l=2 max magnitude:", float(jnp.max(abs(output[cue.SO3(2)]))))

IrrepsNormalize(eps=..., scale=..., skip_scalars=...) uses squared component norms averaged over multiplicities. MLP(layer_sizes=..., activation=..., dtype=..., rngs=...) uses scalar features and fan-in normalization. CUDA indexed-linear species counts must match grouped rows and sum to batch size.

Dense RepArray interface

Use descriptor representation metadata. A single output returns one RepArray: .array holds data; .reps maps axes to representations.

e = cue.descriptors.fully_connected_tensor_product(
    cue.Irreps("SO3", "2x0 + 1"), cue.Irreps("SO3", "0 + 1"),
    cue.Irreps("SO3", "0 + 2x1"),
)
keys = jax.random.split(jax.random.key(11), len(e.inputs))
inputs = [cuex.randn(key, rep, (2,), dtype) for key, rep in zip(keys, e.inputs)]
output = cuex.equivariant_polynomial(e, inputs, method="naive")
np.testing.assert_allclose(output.array, e(*(np.asarray(v.array) for v in inputs))[0],
                           rtol=1e-12, atol=1e-12)
print("RepArray shape:", output.shape, "devices:", output.array.devices())

CUDA preparation and message passing

Prepare compatible dense descriptors with e.squeeze_modes().flatten_coefficient_modes() before equivariant_polynomial(..., method="uniform_1d"). Segment shapes must be uniform within each operand. Splitting by irrep still requires layout conversion.

ir_dict.segmented_polynomial_uniform_1d(poly, inputs, outputs, input_indices=..., output_indices=...) needs leaves shaped (..., op.num_segments, *op.segment_shape), distinct from (..., mul, ir.dim). For _ir_dict operands, transpose/flatten as in split_naive, then reshape using segment metadata. Match leaf order to operand order and output specs to segment suffixes. Reverse conversions before IrrepsLinear.

GNN weights belong to edges, features to nodes, harmonics to edge directions. Gather by senders, scatter by receivers, and size output by destination count. Initialize learned weights/edge networks once in the NNX initializer; split RNG keys for independent inputs. Source examples: cuequivariance_jax/examples/nequip_nnx.py and mace_nnx.py.

Troubleshooting

Symptom	Check
`ensurepip` unavailable	Use a new unique directory, `--without-pip`, and host pip `--python`; preserve partial environments.
Missing CUDA kernels on CPU	Select `naive` or ordinary NNX layers.
Float64 truncated / numerical mismatch	Check early x64, identical operands, dtype/normalization, layout, and operand order.
Dictionary mismatch	Regroup data with metadata; transpose per block. `dict_to_flat` emits `mul_ir`.
Missing output / bad broadcast	Specify every output; broadcast leading axes including weights.
Missing API / unexpected keyword	Inspect that installed symbol's signature/source and version.

Locate source via cuex.__file__; relevant files are ir_dict.py, nnx.py, equivariant_polynomial.py, rep_array/rep_array_.py, and segmented_polynomials/segmented_polynomial*.py. API docs: https://docs.nvidia.com/cuda/cuequivariance/.

Files

ATTRIBUTION.mdAdded by AutoGineer934 B
BENCHMARK.md10.2 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md18.9 KiB
evals/challenges.json6.3 KiB
evals/config.yml2.7 KiB
evals/evals.json6.2 KiB
scripts/requirements.txt218 B
skill-card.md6.5 KiB
skill.oms.sig5.2 KiB

LICENSE

Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.


                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.

      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.

      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).

      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. For the purposes
      of this License, Derivative Works shall not include works that remain
      separable from, or merely link (or bind by name) to the interfaces of,
      the Work and Derivative Works thereof.

      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."

      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.

   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.

   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.

   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and

      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and

      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.

      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.

   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.

   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.

   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.

   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.

   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.

   END OF TERMS AND CONDITIONS

   APPENDIX: How to apply the Apache License to your work.

      To apply the Apache License to your work, attach the following
      boilerplate notice, with the fields enclosed by brackets "[]"
      replaced with your own identifying information. (Don't include
      the brackets!)  The text should be enclosed in the appropriate
      comment syntax for the file format. We also recommend that a
      file or class name and description of purpose be included on the
      same "printed page" as the copyright notice for easier
      identification within third-party archives.

   Copyright [yyyy] [name of copyright owner]

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Created by

Mario Geiger

Publisher

NVIDIA (Official)

Distributed by

NVIDIA

Source

NVIDIA/skills · skills/bionemo-cuequivariance-jax · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-09

AutoGineer is not affiliated with or endorsed by Mario Geiger or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download cuequivariance-jax.zip