Skip to contents

stablehlo

Package website: release | dev

The {stablehlo} R package provides a functional API to create stableHLO programs. These programs can be executed using the R package pjrt.

Installation

From GitHub:

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

You can also install from r-universe, by adding the code below to your .Rprofile.

options(repos = c(
  rxla = "https://r-xla.r-universe.dev",
  CRAN = "https://cloud.r-project.org/"
))

Quickstart

Below, we create a function that takes two input arguments x and y of type f32 and shape 2x2 and adds them. Passing func to hlo_input() is optional, because it will automatically use the last function created with hlo_func().

library(stablehlo)
func <- hlo_func("main")
func
#> func.func @main () ->  {
#> 
#> }
x <- hlo_input("x", "f32", shape = c(2, 2), func = func)
x
#> Variable %x in:
#> func.func @main (%x: tensor<2x2xf32>) ->  {
#> 
#> }
y <- hlo_input("y", "f32", shape = c(2, 2), func = func)
y
#> Variable %y in:
#> func.func @main (%x: tensor<2x2xf32>, %y: tensor<2x2xf32>) ->  {
#> 
#> }
z <- hlo_add(x, y)
z
#> Variable %0 in:
#> func.func @main (%x: tensor<2x2xf32>, %y: tensor<2x2xf32>) ->  {
#> %0 = stablehlo.add %x, %y : tensor<2x2xf32>
#> }
f <- hlo_return(z)
identical(f, func)
#> [1] TRUE
f
#> func.func @main (%x: tensor<2x2xf32>, %y: tensor<2x2xf32>) -> tensor<2x2xf32> {
#> %0 = stablehlo.add %x, %y : tensor<2x2xf32>
#> return %0 : tensor<2x2xf32>
#> }

Main Features

  • Implements type inference to easily build up stableHLO programs.
  • Simple, functional API.
  • Clear error messages.
  • Easy installation because the implementation is in R and does not depend on the stableHLO C++ builder, which depends on LLVM and MLIR.

Important notes

  • stableHLO uses 0-based indexing. Wherever operations take dimension indices (e.g., axes, start indices, permutation dimensions), use 0-based values. This differs from R’s 1-based indexing.
  • We try to use "bool" for the BooleanType. However, in the generated stableHLO representation this will show up as "i1". Think of these as interchangeable.

Limitations

The R package is a partial implementation of the stableHLO specification. At least initially, it will:

  • only support a subset of the available operations, see this issue for an overview.
  • not support all data types, e.g. quantized types and complex numbers are not supported.
  • not support shape dynamism.

Acknowledgments