Create a PJRT Buffer — pjrt

Create a PJRT Buffer from an R object. Any numeric PJRT buffer is an array and 0-dimensional arrays are used as scalars. pjrt_buffer will create a array with dimensions (1) for a vector of length 1, while pjrt_scalar will create a 0-dimensional array for an R vector of length 1.

To create an empty buffer (at least one dimension must be 0), use pjrt_empty.

Important: Uploading a numeric vector at an integer element type rejects any value that type cannot hold: one outside its range, or a missing value, is an error rather than a wrapped or clamped result. A fractional value is not an error – it truncates toward zero, as as.integer() does – and the range is checked on the truncated value, so 255.7 still fits "ui8".

A missing value is not rejected where R's own NA and the element type already share a bit pattern and the vector travels zero-copy: an NA_integer_ at "i32" arrives as INT_MIN, and a bit64::integer64 NA at "i64" arrives as INT64_MIN. Both warn, and as_array() warns about them again on the way back, its check argument defaulting to "warn". At every other integer element type, and at "bool", a missing value is an error – including an integer64 NA at "ui64", where those same bits are the ordinary value 2^63.

At a floating-point element type a missing value is neither rejected nor warned about: it becomes NaN, from an NA_real_ and an NA_integer_ alike, as as.double() would give.

No other checks are performed when creating the buffer – a double too large for "f32", for instance, still becomes Inf.

pjrt_empty() allocates a buffer of the given shape and dtype with unspecified contents. The bytes should be treated as uninitialized — read them only after they have been written to (e.g. as a donated output of pjrt_execute()). Shapes with at least one zero-sized dimension are supported as a degenerate case (the buffer holds zero elements).

Usage

pjrt_buffer(data, dtype = NULL, device = NULL, shape = NULL, ...)

pjrt_scalar(data, dtype = NULL, device = NULL, ...)

pjrt_empty(dtype, shape, device = NULL)

Arguments

data

(any)
Data to convert to a PJRTBuffer.

dtype

(NULL | character(1) | DataType)
The type of the buffer. Currently supported types are:

"bool": boolean ("pred" and "i1" are accepted as aliases)
"{s,u}{8,16,32,64}": Signed and unsigned integer (for integer or double data).
"f{32,64}": Floating point (for double or integer data).
"bf16": bfloat16 (for double or integer data). Values are rounded to the nearest bfloat16 value, ties to even; the largest finite value is about 3.3895e38, and magnitudes at or above the overflow midpoint (about 3.3962e38) round to Inf. as_array() returns the exactly representable values as double.
"f16": IEEE 754 binary16 (for double or integer data). Values are rounded to the nearest binary16 value, ties to even; the largest finite value is 65504, and magnitudes at or above the overflow midpoint 65520 round to Inf. as_array() returns the exactly representable values as double. The default (NULL) depends on the method:
logical -> "bool"
integer -> "i32"
double -> "f32"
raw -> must be supplied

A double at an integer dtype is truncated toward zero, like as.integer() but without its 32-bit intermediate, so pjrt_buffer(2^40, dtype = "i64") stores 1099511627776 rather than overflowing. A value the dtype cannot hold is an error rather than a wrapped or clamped result, and the range is tested after truncation, so 255.7 still fits "ui8".

device

(NULL | PJRTDevice | character(1))
A PJRTDevice object or the name of the platform to use ("cpu", "cuda", ...), in which case the first device for that platform is used. The default is to use the CPU platform, but this can be configured via the PJRT_PLATFORM environment variable. A value the target dtype cannot hold is rejected whatever this is set to, so the flag only governs missing values.

shape

(NULL | integer())
The dimensions of the buffer. The default (NULL) is to infer them from the data if possible. The default (NULL) depends on the method.

...

(any)
Additional arguments. For raw types, this includes:

row_major: Whether to read the data in row-major format or column-major format. R uses column-major format.

Value

PJRTBuffer

Extractors

platform() -> character(1): for the platform name of the buffer ("cpu", "cuda", ...).
device() -> PJRTDevice: for the device of the buffer (also includes device number)
elt_type() -> PJRTElementType: for the element type of the buffer.
shape() -> integer(): for the shape of the buffer.

Converters

as_array() -> array | vector: for converting back to R (vector is only used for shape integer()).
as_raw() -> raw for a raw vector.

Reading and Writing

safetensors::safe_save_file for writing to a safetensors file.
safetensors::safe_load_file for reading from a safetensors file.

A buffer can also be serialized like any R object, e.g. with saveRDS(): its data is copied to the host, and unserializing it uploads the data to the default device (the first device of the PJRT_PLATFORM platform, "cpu" unless set) – not to the device it was saved from. A buffer saved on a GPU therefore also loads on a machine without one; use copy_buffer() to move it afterwards. A buffer referenced several times in one serialized object is written once, and loads as one buffer. A loaded buffer is a NULL external pointer until its first use (e.g. by as_array(), shape() or printing), which fills in its address. This requires serialization format version 3 (R's default); a buffer serialized under version 2 loads without its data.

Scalars

When calling this function on a vector of length 1, the resulting shape is 1L. To create a 0-dimensional buffer, use pjrt_scalar where the resulting shape is integer().

Examples

# Create a buffer from a numeric vector
buf <- pjrt_buffer(c(1, 2, 3, 4))
buf
#> PJRTBuffer 
#>  1
#>  2
#>  3
#>  4
#> [ CPUf32{4} ] 

# Create a buffer from a matrix
mat <- matrix(1:6, nrow = 2)
buf <- pjrt_buffer(mat)
buf
#> PJRTBuffer 
#>  1 3 5
#>  2 4 6
#> [ CPUi32{2x3} ] 

# Create an integer buffer from an array
arr <- array(1:8, dim = c(2, 2, 2))
buf <- pjrt_buffer(arr)
# Create a scalar (0-dimensional array)
scalar <- pjrt_scalar(42, dtype = "f32")
scalar
#> PJRTBuffer 
#>  42
#> [ CPUf32{} ] 
# Allocate an uninitialized 2x3 f32 buffer (contents are unspecified)
empty <- pjrt_empty(dtype = "f32", shape = c(2, 3))
empty
#> PJRTBuffer 
#>  0 0 0
#>  0 0 0
#> [ CPUf32{2x3} ]