Execute a PJRT program

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)

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.

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