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 just-in-time (JIT) compilation for very fast execution and reverse-mode automatic differentiation. Programs can run on CPU and NVIDIA GPU.
Installation
install.packages("anvl", repos = c("https://r-xla.r-universe.dev", getOption("repos")))Afterwards, install the additional dependencies:
anvl::install_anvl()See the installation guide for more details, including prebuilt Docker images.
Why anvl
{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 core ideas:
- Compilation. {anvl} converts R functions into an optimized program via XLA – the same compiler that powers JAX and TensorFlow. Due to the compilation step, 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.
Moreover, the package is designed to be extensible. As the package is written in R, new primitives and transformations can be added without needing a lower-level language.
Usage
We define an R function operating on AnvlArrays, which is 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) {
nv_mul(a, x) + b
}
a <- nv_scalar(1)
b <- nv_scalar(2)
x <- nv_scalar(3)
# Eager mode
f(a, b, x)
#> AnvlArray
#> 5
#> [ CPUf32{} ]
# JIT mode
f_jit <- jit(f)
f_jit(a, b, x)
#> AnvlArray
#> 5
#> [ CPUf32{} ]Through function transformation, 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{} ]By default, floats use f32 precision in {anvl} because of its speed advantage on GPUs. Changing this default to double-precision is possible via the anvl.default_dtypes option:
with_default_dtypes(c(float = "f64"), {
nv_add(pi, 1)
})
#> AnvlArray
#> 4.1416
#> [ CPUf64{} ]For more complex examples, such as implementing a Gaussian Process, see the package website.