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 "pred", 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:

"pred": predicate (i.e. a boolean)
"{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). The default (NULL) depends on the method:
logical -> "pred"
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.

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 
#>  -3.1738e+38  3.0651e-41 -2.8162e+38
#>   3.0651e-41 -2.5557e+38  3.0651e-41
#> [ CPUf32{2x3} ]