numba-enzyme

Differentiate Numba-compiled Python functions with the standalone Enzyme LLVM plugin.

Developed by Yousof Mardoukhi on behalf of Qruise GmbH.

Install

pip install numba-enzyme

Wheels bundle their own copy of clang/opt/llvm-link, the Enzyme LLVM pass plugin, and the C-runtime pieces needed to link the compiled kernels — nothing beyond pip install is required on the target machine.

Supported platforms
OS / architecture Linux x86_64 only
glibc ≥ 2.39 (e.g. Ubuntu 24.04+, Debian 13+, Fedora 39+)
Python CPython 3.11, 3.12, 3.13

Not supported: macOS, Windows, musl-based distros (Alpine), ARM/aarch64, i686, PyPy, or free-threaded builds.

Usage

Annotate each parameter and the return value with a type from numba_enzyme.types, then use grad, jvp, or the @differentiable decorator:

import math
from numba_enzyme.core import grad, jvp, differentiable
from numba_enzyme.types import Float64

def f(x: Float64, y: Float64) -> Float64:
    return x * y + math.cos(x * y)

grad(f)(1.0, 2.0)              # -> (df/dx, df/dy)
jvp(f)((1.0, 2.0), (1.0, 0.0))  # -> directional derivative along (1.0, 0.0)

@differentiable
def g(x: Float64, y: Float64) -> Float64:
    return x * y + math.cos(x * y)

g(1.0, 2.0)          # calls the original Python function directly
g.grad(1.0, 2.0)      # reverse-mode gradient, built lazily on first access
g.jvp((1.0, 2.0), (1.0, 0.0))  # forward-mode JVP

grad returns the gradient with respect to every argument, as a tuple. jvp takes a tuple of primal values and a tuple of tangent (seed) values and returns a single float — the Jacobian-vector product.

Scope

Currently supports pure real scalar math (+, *, **, math.sin/cos/exp, ...) for functions of the shape f(x0: T0, ..., xN: TN) -> TR, both forward- and reverse-mode. Arrays/reductions, np.linalg/BLAS, and complex128 are not supported yet.

License

Apache License 2.0, © Qruise GmbH — see LICENSE.

Project links

Key dates

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

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for ymardoukhi from gravatar.com ymardoukhi

Credits

Author: Yousof Mardoukhi

License expression

Apache-2.0
View SPDX License List

Requires

Python >=3.11

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.

numba_enzyme-0.1.2-cp313-cp313-manylinux_2_39_x86_64.whl (73.4 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.39+ x86-64

numba_enzyme-0.1.2-cp312-cp312-manylinux_2_39_x86_64.whl (73.4 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.39+ x86-64

numba_enzyme-0.1.2-cp311-cp311-manylinux_2_39_x86_64.whl (73.4 MB view details)

Uploaded CPython 3.11manylinux: glibc 2.39+ x86-64

File details

Details for the file numba_enzyme-0.1.2-cp313-cp313-manylinux_2_39_x86_64.whl.

File metadata

File hashes

Hashes for numba_enzyme-0.1.2-cp313-cp313-manylinux_2_39_x86_64.whl
Algorithm Hash digest
SHA256 cf108f1c03b8f862d09b9fdec2c568e463ca4d80cef434372e74f96540d6ede8
MD5 e3951ea7bbb45ffef5734edf8c687ce2
BLAKE2b-256 a2c74e5f7619acb2857192de5d9380680e5d9f23401bbb8f2acc03f68de44e53

See more details on using hashes here.

Provenance

The following attestation bundles were made for numba_enzyme-0.1.2-cp313-cp313-manylinux_2_39_x86_64.whl:

Publisher: wheels.yml on ymardoukhi/numba-enzyme

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file numba_enzyme-0.1.2-cp312-cp312-manylinux_2_39_x86_64.whl.

File metadata

File hashes

Hashes for numba_enzyme-0.1.2-cp312-cp312-manylinux_2_39_x86_64.whl
Algorithm Hash digest
SHA256 5f9cb46eda41e4de22965124c7f0f85d323b009ad519c1251346fbb28905c2f1
MD5 1d45ff7c44f66b66748ecffd89fc8916
BLAKE2b-256 0e695822bba70adabf866ef83c3ad2a2fd46d1141e149a9445cf701d88f089fc

See more details on using hashes here.

Provenance

The following attestation bundles were made for numba_enzyme-0.1.2-cp312-cp312-manylinux_2_39_x86_64.whl:

Publisher: wheels.yml on ymardoukhi/numba-enzyme

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file numba_enzyme-0.1.2-cp311-cp311-manylinux_2_39_x86_64.whl.

File metadata

File hashes

Hashes for numba_enzyme-0.1.2-cp311-cp311-manylinux_2_39_x86_64.whl
Algorithm Hash digest
SHA256 fd67ce71f04e0ff5930215a2bf142ba37af9b728a395c530f99e9b91175caefe
MD5 79948be5f598c1a87e3e21921ae79127
BLAKE2b-256 7e18ef9f8629534409cedb64f6a46f6409127bdeab73693935d359a0866383e6

See more details on using hashes here.

Provenance

The following attestation bundles were made for numba_enzyme-0.1.2-cp311-cp311-manylinux_2_39_x86_64.whl:

Publisher: wheels.yml on ymardoukhi/numba-enzyme

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

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