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: No checks are performed when creating the buffer, so you need to ensure that the data fits the selected element type (e.g., to prevent buffer overflow) and that no NA values are present.

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))
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 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

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.

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().