cuequivariance-ops-torch

Introduction

cuequivariance_ops_torch provides CUDA kernels for the cuEquivariance project's PyTorch components. As such, it contains pytorch bindings to optimized kernels that cuEquivariance's operations map down to. In general, we advice that you access those kernels through cuEquivariance, but you may also find them useful on their own.

Installation

Please install using either pip install cuequivariance-ops-torch-cu12 or pip install cuequivariance-ops-torch-cu13 (depending on the CUDA toolkit you wish to use).

Documentation

For detailed usage information of the kernels, please refer to the doc-strings in their respective modules. For higher-level documentations, refer to cuEquivariance.

Pairformer inference primitives

The Torch package also exposes lower-level primitives for integrating Pairformer inference graphs:

  • attention_pair_bias_mask_from_normalized projects an already-normalized pair representation and applies an arbitrary Boolean [B,U,V] mask. The implementation is shape-generic; BF16 D=256, H=16, and square sequence lengths 384 and 512 are the performance-validated configurations. FP16 and BF16 inputs use native low-precision tensor-core products with FP32 accumulation; FP32 inputs retain the higher-accuracy TF32x3 path.
  • pairformer_combined_triangle_multiplication combines incoming and outgoing triangle multiplication after caller-provided non-affine LayerNorm with epsilon 1e-5. Its fused direction join applies the affine middle LayerNorm with epsilon 0.03. This operation intentionally fails closed outside BF16 inference with B=1, D=256, sequence length 384 or 512, and an SM100 GPU. The caller guarantees that the second normalized carrier is the separately materialized spatial transpose of the first; the hot path validates metadata and distinct storage but does not compare tensor contents. It returns the base update before any caller-owned peri-LayerNorm, output mask, dropout, or residual addition; model-level dropout must be disabled for inference. Use pack_pairformer_grouped_projection to convert grouped projection weights to the expected payload/gate layout.

Both operations are registered as torch.library custom operators with fake implementations for torch.compile tracing. See their Python docstrings for the complete input contracts.

Usage

You can import the library from python:

import cuequivariance_ops_torch

Kernels are primarily exposed as torch.nn.Module, but also provide a lower-level interface as torch.library operators. Generally, the module is responsible for proper input transformation and initialization, and the operator execute the kernel. This allows you to export models using this operations using torch.export, and running inference on them using TensorRT.

Support and Feedback

Please contact the cuEquivariance developers for any issues you might encounter.

Project links

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Latest release

8 maintainers

PyPI data
Data sourced directly from PyPI's database.
Avatar for fpellegrini75 from gravatar.com fpellegrini75 Avatar for hanyichou from gravatar.com hanyichou Avatar for hsadasiv from gravatar.com hsadasiv
+5 more Avatar for jomitchellnv from gravatar.com jomitchellnv Avatar for kucukben from gravatar.com kucukben Avatar for mariogeiger from gravatar.com mariogeiger Avatar for mhohnerbach_nvidia from gravatar.com mhohnerbach_nvidia Avatar for tingyu_nvidia from gravatar.com tingyu_nvidia

Credits

Author: NVIDIA Corporation

License

Other/Proprietary License (NVIDIA Proprietary Software)

Requires

Python >=3.10

Provides Extra

test plugins

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.7 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (462.5 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.9 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (462.2 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.9 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (462.3 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (468.2 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (463.3 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.9 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (463.1 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 7c47d499147559005a61f4ef9b7293bf167bc8a6b622bf8091578125b3cdea8f
MD5 39fde19ecb27c646685c76e4f8d5dff8
BLAKE2b-256 311191e7d777cd52760be9e2bd40dd77978b9b7469e59b8ecbb8821ee0e23a2c

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 462d096dfd0e0cf5706ac529a171ae254cda2c5b893f6ed29ce4bde42d723a34
MD5 39c67907e10287eac3226d57764fa263
BLAKE2b-256 8e1fb0e07d2204e9ad5b1c7b35c02530b8f3c7b96287880662667d0aab4e16bc

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 bc55167e47062e7738f9e28616f60cbe3ff04208b8c34fbb0b2bab8f3f4dfe36
MD5 560a2a4f8288fa433ae67c5dac092807
BLAKE2b-256 8a947b4c5399fb78f3c1ec111b8e8ef747de54a3e4f8d9296979380dcb1eb2eb

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 2d1b3f9e407757012fc5e707553e58459ed6d4f20246843795513e454bca7927
MD5 e25487420bcb27e863d7cdd4b8942bd1
BLAKE2b-256 221ebcfd1e4de1e86b196a56d7c0a72c40dbfc462706c21c3d1f9fe01a9a29a1

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 797fdad536c7ab295d582b2f20c7aeddbdfb04c997ba98c957fe8a5806edb32b
MD5 8e0ec5bb48d64224f8ab79b1799faf13
BLAKE2b-256 45ee0a3ce778e18fc50b953dd956c9264c5cc709495c7f366cd12de78e335a62

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 648f0ecb0e227fdbefe7bcc36d9bba947bb9bc9ac0bda1e586260e65b8d98631
MD5 78c824b031a5c65b69201e958de3bf65
BLAKE2b-256 e1055d2e91ae7122dfd85ce185e480d949ff1ca5ab43bf4a88664d1ad54b64a1

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 adfd61f228dc13eaf33ac34c7ae59479f3358f3c81e6439e8800b4080d209931
MD5 cd02bbb26f2d3d5efa161d5155d6ecf8
BLAKE2b-256 8952ee3062a139798aa22fc7b54587ca2e19a57568a8dc3fc1bf07bf2766114e

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 69fca0d1cbc9daf3b3a23d5fe9d5e423a2f95d58eaa8327e72ce31fb90fe081e
MD5 40d40d561d9fadc2740bff9b2190bad0
BLAKE2b-256 0c7f5d29ce91b8cd54f2f4d0d747f77b5f8e62acfecc9e0e2f78e9258f8e04ba

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 b698eb19a8b046d5c7752b53178ca00d0bce47cdae556a61fe4f7f720c606632
MD5 1e8e788534bf40f90215a438045f9ac9
BLAKE2b-256 cea67b51e956fa9d389bd3b64ec39ee55682767222e471eadd673994f54abda8

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 3dd20f9fd88b910d4e2b4c067a2c113cca363630baa6483bada7440c9065a147
MD5 0a87a052038e49c21373bfc9b962213e
BLAKE2b-256 37a057223cfd8696e4a77cc910b7eb98cdf3c776876de7e6d25941f0266b83c3

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page