Skip to contents

pjrt: R Interface to PJRT

Source: R/package.R
pjrt-package.Rd

Provides an R interface to PJRT (Pluggable Jit RunTime), which allows you to run XLA or stableHLO programs on a variety of hardware backends including CPU, GPU, and TPU.

Environment Variables

Configuration options provided by XLA

XLA provides various configuration options, but their documentation is scattered across various websites. The options include:

  • TF_CPP_MIN_LOG_LEVEL: Logging level for PJRT C++ API:

    • 0: shows info, warnings and errors

    • 1: shows warnings and errors

    • 2: shows errors

    • 3: shows nothing

  • XLA_FLAGS: See the openxla website for more information. Among these, the following flags control HLO dumping and are used by inspect_hlo() to inspect the XLA IR:

    • --xla_dump_to=<dir>: Directory XLA writes its HLO dump files into.

    • --xla_dump_hlo_as_text: Dump the HLO modules as text (before and after optimizations).

    XLA reads XLA_FLAGS only once, before the first compilation, so set it at the start of a fresh R session, e.g. Sys.setenv(XLA_FLAGS = "--xla_dump_to=/tmp/hlo --xla_dump_hlo_as_text").

Configuration options provided by this package

  • PJRT_PLATFORM: Default platform to use, falls back to "cpu".

  • PJRT_PLUGIN_PATH_<PLATFORM>: Path to custom plugin library file for a specific platform (e.g., PJRT_PLUGIN_PATH_CPU, PJRT_PLUGIN_PATH_CUDA, PJRT_PLUGIN_PATH_METAL). If set, the package will use this path instead of downloading the plugin.

  • PJRT_PLUGIN_URL_<PLATFORM>: URL to download plugin from for a specific platform (e.g., PJRT_PLUGIN_URL_CPU, PJRT_PLUGIN_URL_CUDA, PJRT_PLUGIN_URL_METAL). If set, overrides the default plugin download URL.

  • PJRT_INSTALL: Controls whether plugins may be downloaded automatically. Set this to "1" to always download without asking (e.g. in CI, scripts, or Docker builds), or to "0" to never download (the call errors with instructions instead). When unset, the package asks for confirmation in an interactive session and errors in a non-interactive one, so a script never triggers a surprise download.

  • PJRT_ZML_ARTIFACT_VERSION: Version of ZML artifacts to download. Only used when downloading plugins from zml/pjrt-artifacts.

  • PJRT_CPU_DEVICE_COUNT: The number of CPU devices to use. Defaults to 1. This is primarily intended for testing purposes.

  • PJRT_CUDA_R_PACKAGE: Name of the R package providing CUDA libraries. Defaults to the value of cuda12.8. Set this to use a different CUDA toolkit package, but note that other versions may not work with the XLA plugin.

  • PJRT_DEBUG: If set (to any non-empty value), enables verbose debug output via cli::cli_inform().

Third-Party Licenses

The pjrt package itself is MIT-licensed. The CUDA backend dynamically loads NVIDIA software which is not bundled with pjrt, but downloaded from NVIDIA's official redistributable channels by the CUDA toolkit R package (e.g. cuda12.8) at install time. Its use is governed by the NVIDIA CUDA Toolkit EULA, with the exception of cuDNN, which is covered by the NVIDIA cuDNN SLA, and NCCL, which is covered by its own license. By installing or using the CUDA backend you accept those terms.

See also

Useful links:

Author

Maintainer: Sebastian Fischer seb.fischer@tutamail.com (ORCID)

Authors:

Other contributors:

  • Daniel James (Author of the Boost.ContainerHash hash_combine copied in src/hash.h (BSL-1.0); see inst/COPYRIGHTS) [copyright holder]

  • The OpenXLA Authors (Authors of the PJRT C API and XLA FFI headers in inst/include and the XLA protobuf definitions in inst/proto (Apache-2.0); see inst/COPYRIGHTS) [copyright holder]