drjit

What it is and what it does

Dr.Jit is a specialized just-in-time compiler designed to handle large, sparse computation graphs typical of differentiable rendering and similar embarrassingly parallel workloads. Unlike general machine learning frameworks, it records arithmetic operations into a computation graph, then JIT-compiles that graph into fused kernels targeting GPUs (Metal on macOS, CUDA elsewhere) or CPUs with vector instruction sets (AVX512, NEON) via LLVM. It optionally computes derivatives using forward or reverse-mode automatic differentiation, with both tracing and differentiation producing specialized code.

The package originated as the numerical foundation of Mitsuba 3, a differentiable Monte Carlo renderer, but is general-purpose and can be used without JIT compilation as a header-only vector library. It supports both C++17 and Python, allowing code to be developed in either language or both simultaneously, with joint tracing and differentiation across language boundaries. Dr.Jit includes a mathematical library with transcendental functions and types like vectors, matrices, complex numbers, and quaternions, and handles custom data structures, side effects, and polymorphism.

Use it for:

• Differentiable rendering pipelines where computation graphs contain millions of elementary operations that would overwhelm ML framework compilers.
• Physics simulations requiring automatic differentiation and GPU acceleration without the overhead of general-purpose ML frameworks.
• Embarrassingly parallel numerical computations that benefit from vectorization and JIT compilation to specialized kernels.
• Research combining Python prototyping with C++ performance-critical code, jointly traced and differentiated.
• Monte Carlo methods and inverse problems where reverse-mode AD is essential for gradient computation.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Dr.Jit is a just-in-time compiler for differentiable and ordinary computation that traces arithmetic operations into computation graphs, JIT-compiles them to GPU or CPU kernels, and optionally computes derivatives via automatic differentiation.

Yes, if you are working on differentiable rendering, physics simulation, or large sparse computation graphs where ML frameworks struggle. The active maintenance, permissive license, and lack of known vulnerabilities support adoption. Medium install friction is acceptable given the specialized use case and broad platform coverage. Not recommended for general machine learning; use JAX, PyTorch, or TensorFlow instead.

Install

drjit on PyPI

pip install drjit

uv add drjit

poetry add drjit

Installing drjit

Before you install

Medium install friction due to compiled wheels for multiple Python versions and architectures (cp310–cp313, arm64/x86_64/Windows). Active maintenance with a release 7 days old and recent commits; no known vulnerabilities.

License in practice

Permissive BSD license allows commercial and private use with minimal restrictions, making it suitable for research and production deployments.

Quickstart

pip install drjit
import drjit as dr
# Create and trace a computation
x = dr.arange(dr.Float, 10)
y = x + 1

Requires Python >=3.8; GPU support (Metal on macOS, CUDA elsewhere) is optional but recommended for performance.

Verify before relying

• Specific performance characteristics compared to JAX, TensorFlow, or PyTorch for large sparse computation graphs.
• Whether the package is actively maintained beyond the Mitsuba 3 renderer project.
• Availability and quality of documentation beyond readthedocs.io.
• Practical limits on computation graph size and complexity.

Package facts

Evidence: drjit-1.5.0-cp310-cp310-macosx_11_0_arm64.whl; drjit-1.5.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl; drjit-1.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; drjit-1.5.0-cp310-cp310-win_amd64.whl; drjit-1.5.0-cp311-cp311-macosx_11_0_arm64.whl; drjit-1.5.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl; drjit-1.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; drjit-1.5.0-cp311-cp311-win_amd64.whl; drjit-1.5.0-cp312-abi3-macosx_11_0_arm64.whl; drjit-1.5.0-cp312-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl; drjit-1.5.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; drjit-1.5.0-cp312-abi3-win_amd64.whl; drjit-1.5.0-cp312-cp312-macosx_11_0_arm64.whl; drjit-1.5.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl; drjit-1.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; drjit-1.5.0-cp312-cp312-win_amd64.whl; drjit-1.5.0-cp313-cp313-macosx_11_0_arm64.whl; drjit-1.5.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl; drjit-1.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; drjit-1.5.0-cp313-cp313-win_amd64.whl

Tags

More Scientific/Engineering packages