liveBook · Manning

Jax.devices()

As explained in Deep Learning with JAX

A JAX function that returns the full set of devices available across processes.

Returns the available global devices, across all processes. In single-host setups, that list usually matches the local devices list; the first entry is the default device for uncommitted arrays. section 3.2.3 section 5.1.1

Commonly used to:

  • inspect which CPU, GPU, or TPU devices JAX can see. section 3.2.3
  • pick a specific device for placement or JIT execution. section 5.1.1
  • compare local versus global device counts in multi-process settings. section 3.2.3

Basic device listing

`jax.devices()` can return a list like `[gpu(id=0)]`, while `jax.devices('cpu')` can return `[CpuDevice(id=0)]`. That makes it easy to see what hardware is available before placing data or compiling a function for a target device. section 3.2.3

Pairs with `jax.device_put()` and `jax.jit(device=...)` when you need precise placement or execution control. It also sits alongside `jax.local_devices()`, `jax.device_count()`, and `jax.local_device_count()` for hardware introspection. section 3.2.3 section 5.1.1

How Deep Learning with JAX covers it

How deep it goes
Explains global versus local devices, backend filtering, and default-device behavior without going into multi-process details.
Examples
Uses device listings such as `[gpu(id=0)]` and `[CpuDevice(id=0)]`, plus comparisons against `jax.local_devices()` and `jax.device_count('gpu')`.
What it assumes
Assumes familiarity with basic JAX arrays and device placement, and introduces global devices and backend names from zero.
What you can do after
Apply it to inspect available hardware and choose a specific device for placement or compilation.
How it connects
Links to `jax.local_devices()`, `jax.device_count()`, `jax.device_put()`, `jax.jit()`, and `jax.process_index()`.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.