StableHLO is an operation set for high-level operations (HLO) in machine learning (ML) models.

It's the portability layer between ML frameworks (targeted for GoMLX, but could be used for others) and ML compilers. It allows for easy support for different vendors, by coupling with XLA's PJRT (*) API for executing StableHLO programs. So many different GPUs and TPUs are supported.

(*) PJRT, which stands for Pluggable JIT Runtime, is an API in the context of XLA (Accelerated Linear Algebra) that provides a unified, cross-platform interface for interacting with different hardware accelerators. StableHLO is the device-independent language to specify the computation, and it also includes APIs to handle buffer (the data) management and optionally distributed execution.

See:

• StableHLO specification
• GoMLX: a Go ML framework that supports an XLA (StableHLO+PJRT) backend to efficiently run (or train) ML programs.
• GoPJRT: a Go wrapper for PJRT C API, capable of executing StableHLO programs, for a lower level API.

The tests in tests/gopjrt/gopjrt_test.go should serve as simple examples of each operation.

Notice that stablehlo is a low-level API, usually used to build higher-level frameworks (an ML framework like GoMLX, maybe an image manipulation library that uses accelerators like GPUs, some scientific library, etc.), so it's deliberately verbose and requires boilerplate (error handling) everywhere. It sacrifices ergonomics for performance, consistency and stability.

See another example of stablehlo and GoPJRT (to execute the generate StableHLO program) in Mandelbrot mandelbrot.ipynb notebook. It includes some sample StableHLO code, if you are curious.

Most operations are already implemented. See the list of supported operations (the ones not implemented are in the bottom of the list).

If you need a specific operation, please open an issue.

See also the CHANGELOG.

• Reference StableHLO documentation here.
• RFC: Dynamism 101

In the first version we aim at supporting only unbounded dynamism using shape polymorphism: where axes dimensions are not defined and has no bounds, and where PJRT will be able to dynamically re-instantiate and re-compile the program to a new shape (or re-use a cache).

Other types of dynamism:

• Unranked dynamism: rank unknown and compile time. Not supported.
• Data-dependent dynamism: for data-dependent dynamic ops. For instance, if a function returns the indices of all non-zero elements. There is little support for this, so we are not support it yet.

With the following advantages:

• XlaBuilder has become a second-class citizen, so to say, within OpenXLA. And things are moving towards the "MLIR builder" (MLIR is the generic ML Intermediary Language, of which StableHLO is a specialization/extension). So we will eventually need a newer "builder" for Gopjrt.
• Since PJRT takes StableHLO in plain text format, we can write this entirely in Go, not requiring any extra C/C++ library build.
  • PJRT itself is a C library, but with a relatively small API surface, and for which there are prebuilt distributions available (for Jax). So we can get away without having to manage Bazel issues.
  • The goal is to eventually not require a C compiler to compile gopjrt, and instead use ebitengine/purego do dynamically load PJRT.
  • There are PJRT for different platforms. If we don't need to compile XlaBuilder for them, it makes it more plausible to support them.

The disadvantages:

• XlaBuilder provided "shape inference." So if I say Add(a, b) the XlaBuilder would tell how to broadcast a and b, and the resulting shape. When we build the StableHLO we have to re-implement this shape inference, not only for the GoPJRT users, but also because the StableHLO language requires the inputs and outputs shapes to be specified in every statement.
  • This is not a disadvantage for the user of this library, since stablehlo does that for you, but it's more work for the library maintainers.
• This means more maintenance: any updates in the language specification or new ops need to have their shape inference updated accordingly.

The same code is also used by GoMLX SimpleGo engine (github.com/gomlx/gomlx/backends/simplego), but we didn't want to create a dependency in either direction: users of Gopjrt may not be interested in GoMLX, and users of GoMLX that don't use the XLA backend wouldn't want a dependency to Gopjrt.

Also, some operations have slightly different nuances.