Constructs the PJRT backend, which stores array data in PJRT buffers (via
pjrt::pjrt_buffer()) and compiles jitted functions to XLA executables
via stablehlo() and pjrt::pjrt_compile(). This is the default
backend.
Value
(AnvlBackend)
With subclass "AnvlBackendPjrt".
Data representation
An AnvlArray with backend = "pjrt" wraps a pjrt::pjrt_buffer()
stored in the $data field. The buffer owns the memory holding the array
values and may live on any device supported by PJRT (CPU, CUDA, Metal,
...). Calling as_array() transfers the buffer contents back to an R
array; calling nv_array() on an R object uploads it to the requested
device.
Each AnvlArray therefore has an associated device, queryable via
device(). A device is a pjrt::as_pjrt_device() object (e.g. the
platform "cpu" or "cuda", optionally with an index such as "cuda:1").
When device is NULL in nv_array() or the jit() wrapper, the
device defaults to default_device(), or is inferred from the existing
inputs of a jitted call.
Operations require all inputs to live on the same device.
Supported data types
bool; the signed integers i8, i16, i32 and i64; the unsigned
integers ui8, ui16, ui32 and ui64; and the floats f32 and f64.
An R double materializes at f32 on this backend and an R integer at i32
unless the defaults say otherwise (see default_dtypes()).
Floating-point behavior
Subnormal floating-point values may be preserved when stored in an array and read back into R, yet treated as zero in calculations. On CPUs, XLA enables a mode that replaces subnormal inputs and results with zero. The exact behavior depends on the platform, backend, and operation.
See vignette("gotchas", package = "anvl") for an explanation and examples.
PJRT JIT arguments
donate(character(), defaultcharacter()): names of arguments whose underlying buffers may be donated to (i.e., reused/consumed by) the compiled XLA executable. Donated buffers must not be used again by the caller after the call; this can reduce memory usage and copies for large inputs. Must not overlap withstatic.