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:
Newer release available (0.11.1)

5 maintainers

PyPI data
Data sourced directly from PyPI's database.
Avatar for fpellegrini75 from gravatar.com fpellegrini75 Avatar for hsadasiv from gravatar.com hsadasiv Avatar for jomitchellnv from gravatar.com jomitchellnv
+2 more Avatar for mariogeiger from gravatar.com mariogeiger 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_cu13-0.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (472.0 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (467.7 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (472.1 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (467.5 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (472.3 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (467.5 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (472.5 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (468.5 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (472.2 kB view details)

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

cuequivariance_ops_torch_cu13-0.11.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (468.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_cu13-0.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 0e02d49060b6df233a7489e56766c77990e3ec55a2f89c4959c9fd6b62b8a701
MD5 c7053b64c10889344cc68da84dd8e6e3
BLAKE2b-256 17c482bca23464c668e69061234836bb8868b934a25090bda8d66ad4c4b992e2

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 953a6aa19675bf912009896ceb798184ed4f0ec606ab9eca849394d851ea2575
MD5 f684590a13279c9a06967c71eb75e173
BLAKE2b-256 5ab04d5ffbdab0c4f89b21451b37da0bf0a41b28894d6524621c1a989e2072c1

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 7db7e6c5fcc370f3d219f3a90d35b1f609a4b8b728dbfbdfa5530b716adc6401
MD5 76f4b878d568ae5fdbd533a80582840a
BLAKE2b-256 00084bf363ca13f0a0cd74c1b7e357296addf8afe9a98b3286df5ba63d744785

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 eef616622411e5e16a47bf03a9f7ab7421c5d1c7bb6849c0e1f468b097cb0171
MD5 503ba784c8f1dd4c783e40656dd5f3a9
BLAKE2b-256 b26afa7ac3f0bc7ba7790357096ca0a9fb12b9079f4c351c8723edf167d49212

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 29b78a39a585b6de49869099edd9e96bf7ccc3a2467f0f415484418e0dd8b131
MD5 b124ff74399b38a2bd047477603cfa20
BLAKE2b-256 69e1d63e0b3d0788748eb381ea80ba5e4b651973c90450b1c61c149affa166bc

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 637d7639cbbbdc280a71832748e554e0ea41ade6da04ae1f3354d06339f7b2fc
MD5 963f61525811ed990b2a7ed5876029bb
BLAKE2b-256 1236a61211dd7e9f24c74c3abc9b10551a2fc8b01000ae7f5d662fe7f902874f

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 fb8cc8fa7fee089b6ba9effaf2feb9e1a9ebee01fef878ad1fc16050b8534ff1
MD5 532a8b78bfe09d3fed352b35ad401d4e
BLAKE2b-256 8607bb46f697bd665b216b9c2d51136afd0e3517ddd67ed8b3095828d90e9930

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 5ff4fd94a1ca5d5809f167dee50ea8e7370b0aacdb3db917786e201859685dec
MD5 7d2e50909f4b4e64f2a02ebef985f254
BLAKE2b-256 1fd48b297071a4cc124df2e3ca455a8030814367f5f536035222c80480464ae3

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 a1866450d6c39be3e32ee69dd332b0bb9b0a935cf53da9435bea0758030fbc55
MD5 ac26c99c1b3ce814fcce6a7820d93fdb
BLAKE2b-256 6885c3b3a42b7adb9c7234ddc3d2f570c0fa55542099c7ff438df0b0109a9bef

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu13-0.11.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu13-0.11.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 fe1cba4e73422e26b8dbcba5349b43b1949bd0a67c2f5f987fb20c66c7be05bd
MD5 14412ec06f5e4418609473073323511b
BLAKE2b-256 7d80759271c430976ba6607f50b91b3ecebb79a02bc3c435afe33b8bc690aebb

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