Temporarily set the current MLX device or stream — with_device • Rmlx

Use local_device() to temporarily switch devices within the current function.

Usage

with_device(device, code)

local_device(device, .local_envir = parent.frame())

Arguments

device: "gpu", "cpu", or an mlx_stream created via mlx_new_stream().
code: Expression to evaluate while device is active.
.local_envir: Environment to bind the restoration to. Defaults to the calling environment.

Value

The result of evaluating code.

Invisibly returns the previous default device.

See also

mlx.core.default_device

Examples

with_device("cpu", x <- mlx_vector(1:10))

local_device("cpu")
# code here runs on CPU, then the previous default is restored