Skip to contents

Select JIT device from a function argument

Source: R/jit.R
device_arg.Rd

Pass the result to jit()'s device argument to indicate that the device should be read from a formal argument of the function being compiled. At call time, the value of that argument is used to derive the backend via backend() dispatch and is forwarded to the backend-specific JIT as the compilation device.

This is intended for functions that have no dynamic array inputs from which the backend could otherwise be detected (e.g. array constructors like prim_fill() or prim_iota()).

Usage

device_arg(argname)

Arguments

argname

(character(1))
Name of a formal argument of the function passed to jit().

Value

(AnvlDeviceArg)
An object recognized by jit().

See also

jit(), backend()

Examples

f <- function(x) nv_scalar(1, device = x)
g <- jit(f, backend = "auto", device = device_arg("x"))
g(nv_device("cpu", "xla"))
#> AnvlArray
#>  1
#> [ CPUf32{} ]