Skip to contents

pjrt

The {pjrt} package provides an R interface to PJRT (Pretty much Just another RunTime), which allows you to run XLA and stableHLO programs on various hardware backends. These programs are framework and hardware agnostic, which means they can be generated by ML frameworks such as jax, and run by PJRT on a specified backend (CPU, GPU, etc.). For a low-level R interface to create stableHLO programs, see the stablehlo package.

Installation 

pak::pak("r-xla/pjrt")

Quickstart

Below, we create and run a stableHLO program that adds two f32 tensors of shape (2, 2).

library(pjrt)
src <- r"(
func.func @main(
  %x: tensor<2x2xf32>,
  %y: tensor<2x2xf32>
) -> tensor<2x2xf32> {
  %0 = "stablehlo.add"(%x, %y) : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  "func.return"(%0): (tensor<2x2xf32>) -> ()
}
)"
program <- pjrt_program(src, format = "mlir")
program
#> PJRTProgram(format=mlir, code_size=221)
#> 
#> func.func @main(
#>   %x: tensor<2x2xf32>,
#>   %y: tensor<2x2xf32>
#> ) -> tensor<2x2xf32> {
#> ...
executable <- pjrt_compile(program, client = "cpu")

x <- matrix(as.double(1:4), nrow = 2)
y <- matrix(as.double(5:8), nrow = 2)
x_buffer <- pjrt_buffer(x, "f32")
y_buffer <- pjrt_buffer(y, "f32")

result <- pjrt_execute(executable, x_buffer, y_buffer)

as_array(result)
#>      [,1] [,2]
#> [1,]    6   10
#> [2,]    8   12

Main Features

  • Compile stableHLO programs into hardware-specific executables.
  • Provide a runtime to execute compiled programs.
  • Convert buffers to and from R arrays and vectors.
  • Read and write buffers using the safetensors format.

Platform Support

  • Linux
    • ✅ CPU backend is fully supported.
    • ✅ CUDA (NVIDIA GPU) backend is fully supported.
  • Windows
    • ⚠️ Currently only supported via Windows Subsystem for Linux (WSL2).
  • macOS
    • ✅ CPU backend is supported.
    • ⚠️ Metal (Apple GPU) backend is available but not fully functional.

Acknowledgements

The design of the {pjrt} package was inspired by the gopjrt implementation.