Skip to contents

Execute a PJRT program

Source: R/loaded_executable.R
pjrt_execute.Rd

Execute a PJRT program with the given inputs and execution options. Returns immediately with PJRTBuffer object(s) that may not be ready yet.

Important: Arguments are passed by position and names are ignored.

Inputs can be PJRTBuffer objects, including buffers that are not yet ready. PJRT handles the dependencies internally.

Use await() to block until the result is ready. Use is_ready() to check if execution has completed (non-blocking). Use as_array_async() to chain async buffer-to-host transfer.

Usage

pjrt_execute(
  executable,
  ...,
  execution_options = NULL,
  simplify = TRUE,
  check = TRUE
)

Arguments

executable

(PJRTLoadedExecutable)
A PJRT program

...

(PJRTBuffer)
Inputs to the program. Named are ignored and arguments are passed in order.

execution_options

(PJRTExecuteOptions)
Optional execution options for configuring buffer donation and other settings.

simplify

(logical(1))
If TRUE (default), a single output is returned as a PJRTBuffer. If FALSE, a single output is returned as a list of length 1 containing a PJRTBuffer.

check

(logical(1))
If TRUE (default), validate the arguments (executable type, that all inputs are PJRTBuffers and unnamed, execution options, and simplify). Trusted callers on a hot dispatch path can pass FALSE to skip this validation; the caller is then responsible for guaranteeing that executable, the inputs, and execution_options are valid.

Value

PJRTBuffer | list of PJRTBuffers

See also

await(), is_ready(), as_array_async()

Examples

# Create and compile a simple identity program
src <- r"(
func.func @main(
  \%x: tensor<2x2xf32>,
  \%y: tensor<2x2xf32>
) -> tensor<2x2xf32> {
  \%0 = "stablehlo.add"(\%x, \%y) : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  "func.return"(\%0): (tensor<2x2xf32>) -> ()
}
)"
prog <- pjrt_program(src = src)
exec <- pjrt_compile(prog)
#> Error: -:3:3: error: unexpected character
#> <unknown>:0: error: Failed to parse using StableHLO v1.13.2, this could indicate forward incompatibility, >12w old unsupported plugin, or a portable artifact that needs to be further downgraded.

# Execute with input
x <- pjrt_buffer(c(1.0, 2.0, 3.0, 4.0), shape = c(2, 2), dtype = "f32")
y <- pjrt_buffer(c(5, 6, 7, 8), shape = c(2, 2), dtype = "f32")
pjrt_execute(exec, x, y)
#> Error: object 'exec' not found