LuxDeviceUtils
LuxDeviceUtils.jl is a lightweight package defining rules for transferring data across devices. Most users should directly use Lux.jl instead.
Index
Preferences
# LuxDeviceUtils.gpu_backend! —
Function.
```julia
gpu_backend!() = gpu_backend!("")
gpu_backend!(backend) = gpu_backend!(string(backend))
gpu_backend!(backend::AbstractLuxGPUDevice)
gpu_backend!(backend::String)
```
Creates a `LocalPreferences.toml` file with the desired GPU backend.
If `backend == ""`, then the `gpu_backend` preference is deleted. Otherwise, `backend` is validated to be one of the possible backends and the preference is set to `backend`.
If a new backend is successfully set, then the Julia session must be restarted for the change to take effect.
Data Transfer
# LuxDeviceUtils.cpu_device —
Function.
```julia
cpu_device() -> LuxCPUDevice()
```
Return a `LuxCPUDevice` object which can be used to transfer data to CPU.
# LuxDeviceUtils.gpu_device —
Function.
```julia
gpu_device(; force_gpu_usage::Bool=false) -> AbstractLuxDevice()
```
Selects GPU device based on the following criteria:
1. If `gpu_backend` preference is set and the backend is functional on the system, then that device is selected.
2. Otherwise, an automatic selection algorithm is used. We go over possible device backends in the order specified by `supported_gpu_backends()` and select the first functional backend.
3. If no GPU device is functional and `force_gpu_usage` is `false`, then `cpu_device()` is invoked.
4. If nothing works, an error is thrown.
Miscellaneous
# LuxDeviceUtils.reset_gpu_device! —
Function.
```julia
reset_gpu_device!()
```
Resets the selected GPU device. This is useful when automatic GPU selection needs to be run again.
# LuxDeviceUtils.supported_gpu_backends —
Function.
```julia
supported_gpu_backends() -> Tuple{String, ...}
```
Return a tuple of supported GPU backends.
::: warning
This is not the list of functional backends on the system, but rather backends which `Lux.jl` supports.
:::
::: danger
`Metal.jl` support is **extremely** experimental and most things are not expected to work.
:::