anvl

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.

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.

{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.

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{} ]

To learn more, see the Get Started article.

✓ supported  ·  WSL2: via the Windows Subsystem for Linux  ·  ✗ not supported

• This work is supported by MaRDI.
• The design of this package is inspired by JAX.
• For JIT compilation, we build on the OpenXLA project.
• Most of the compiler binaries are downloaded from zml.