Package website: release | dev
Accelerated array computing and code transformations for R, allowing you to run numerical programs at the speed of light. The package supports JIT compilation for very fast execution and reverse-mode automatic differentiation. Programs can run on CPU and NVIDIA GPU.
install.packages("anvl", repos = c("https://r-xla.r-universe.dev", getOption("repos")))Install CUDA support on linux amd64:
install.packages("cuda12.8", repos = "https://mlverse.r-universe.dev")See the Installation guide for more details, including prebuilt Docker images.
anvl makes numerical R code run fast on CPUs and GPUs, and computes gradients of your functions automatically. It aspires to be for R what JAX is for Python.
There are three ideas:
- Compilation. {anvl} converts R functions into an optimized program via XLA – the same compiler that powers JAX and TensorFlow. Due to the compilation, resulting programs can be faster compared to implementing them in {torch}.
- Function transformation. Programmatically derive new functions
from existing ones. Currently the only available transformation is
reverse-mode automatic differentiation via
gradient(), which returns the derivative of a function as another R function. - Hardware portability. The same code runs on CPU or GPU.
- Extensible. The package is written almost entirely in R; new primitives and transformations can be added without leaving R.
We define an R function operating on AnvlArrays – the primary data
type of {anvl}. It can be executed in either eager mode (each
operation is performed immediately) or jit mode (the whole function is
compiled into a single executable via jit()).
library(anvl)
f <- function(a, b, x) {
a * x + b
}
a <- nv_scalar(1.0, "f32")
b <- nv_scalar(-2.0, "f32")
x <- nv_scalar(3.0, "f32")
# Eager mode
f(a, b, x)
#> AnvlArray
#> 1
#> [ CPUf32{} ]
# JIT mode
f_jit <- jit(f)
f_jit(a, b, x)
#> AnvlArray
#> 1
#> [ CPUf32{} ]Through automatic differentiation, we can also obtain the gradient of the above function.
g_jit <- jit(gradient(f, wrt = c("a", "b")))
g_jit(a, b, x)
#> $a
#> AnvlArray
#> 3
#> [ CPUf32{} ]
#>
#> $b
#> AnvlArray
#> 1
#> [ CPUf32{} ]For more complex examples, such as implementing a Gaussian Process, see the package website.
| Platform | CPU | GPU |
|---|---|---|
| Linux (x86_64) | ✓ | ✓ CUDA |
| Linux (ARM) | ✓ | ✗ |
| Windows | ✓ | ◐ WSL2 only (CUDA) |
| macOS (Apple Silicon) | ✓ | ✗ |
| macOS (Intel) | ✗ | ✗ |
✓ fully supported · ◐ limited support · ✗ not supported
- This work is supported by MaRDI.
- The design of this package was inspired by and borrows from:
- JAX, especially the autodidax tutorial.
- The microjax project.
- For JIT compilation, we leverage the OpenXLA project.