Device¶

tinax.device owns backend-agnostic pre-JAX environment policy that the JAX runtime honors, a summary of the active JAX runtime, and optional per-vendor accelerator detection gated behind the accelerator extras. It does not select hardware or apply out-of-memory heuristics; those policies belong to the application.

from tinax.device import configure_jax, configure_single_chip, device_info, set_visible_cuda

set_visible_cuda([0, 1])          # before importing jax
configure_single_chip(0)          # expose one local TPU chip before importing jax
configure_jax(preallocate=False, matmul_precision="high")

info = device_info()              # imports jax on demand

Detection is explicit per vendor, and each function names the extra that installs its binding:

Function	Binding	Extra
`nvidia_gpus`	`nvidia-ml-py`	`tinax[cuda12]`, `tinax[cuda13]`, `tinax[cuda12-local]`, `tinax[cuda13-local]`
`amd_gpus`	`amdsmi`	`tinax[rocm7-local]`
`intel_gpus`	`pyzes`	`tinax[oneapi]`
`tpus`	`tpu-info`	`tinax[tpu]`

from tinax.device import amd_gpus, intel_gpus, nvidia_gpus, tpus

for gpu in nvidia_gpus():
    print(gpu.backend, gpu.index, gpu.name, gpu.free_bytes, gpu.compute_capability)

for gpu in amd_gpus():
    print(gpu.backend, gpu.index, gpu.name, gpu.free_bytes, gpu.architecture)

for gpu in intel_gpus():
    print(gpu.backend, gpu.index, gpu.name, gpu.free_bytes)

Every GpuInfo uses the same fields. compute_capability is NVIDIA's SM version and architecture is AMD's gfx target; memory fields are None when the vendor binding cannot report them.

set_visible_cuda ¶

set_visible_cuda(indices: Iterable[int]) -> None

Set CUDA_VISIBLE_DEVICES from explicit device indices before JAX initializes.

Parameters:

Name	Type	Description	Default
`indices`	`Iterable[int]`	Non-negative NVIDIA device indices to expose, in order.	required

Raises:

Type	Description
`RuntimeError`	If JAX has already been imported.
`TypeError`	If an index is not an integer, or is a boolean.
`ValueError`	If an index is negative.

configure_single_chip ¶

configure_single_chip(chip: int) -> None

Expose exactly one local TPU chip to a forthcoming JAX import.

Parameters:

Name	Type	Description	Default
`chip`	`int`	Non-negative local TPU chip index to expose.	required

Raises:

Type	Description
`RuntimeError`	If JAX has already been imported.
`TypeError`	If `chip` is not an integer, or is a boolean.
`ValueError`	If `chip` is negative.

configure_jax ¶

configure_jax(
    *,
    preallocate: bool = False,
    matmul_precision: str | None = "high",
) -> None

Set JAX/XLA runtime options before JAX initializes.

Parameters:

Name	Type	Description	Default
`preallocate`	`bool`	Whether XLA may preallocate the device memory pool; must be set before JAX initializes.	`False`
`matmul_precision`	`str \| None`	Default matmul precision applied via `jax.config` once JAX is imported, or `None` to leave it unchanged.	`'high'`

Raises:

Type	Description
`RuntimeError`	If JAX has already been imported.
`TypeError`	If `preallocate` is not a boolean, or `matmul_precision` is neither a string nor `None`.

device_info ¶

device_info() -> DeviceInfo

Summarize the active JAX runtime backend and its visible devices.

Returns:

Name	Type	Description
`A`	`DeviceInfo`	class:`DeviceInfo` describing the default backend, global and local device
	`DeviceInfo`	counts, this process's index and the process count, and the visible device
	`DeviceInfo`	kinds. Importing this module stays inert; JAX is imported only when called.

DeviceInfo `dataclass` ¶

DeviceInfo(
    backend: str,
    device_count: int,
    local_device_count: int,
    process_index: int,
    process_count: int,
    device_kinds: tuple[str, ...],
)

Immutable summary of the active JAX runtime backend and its visible devices.

nvidia_gpus ¶

nvidia_gpus() -> list[GpuInfo]

Probe visible NVIDIA GPUs through nvidia-ml-py (imported as pynvml).

Returns:

Type	Description
`list[GpuInfo]`	The visible NVIDIA GPUs in device-index order, or an empty list when none are
`list[GpuInfo]`	present. Memory fields are byte counts and `compute_capability` is reported.

Raises:

Type	Description
`ImportError`	If `nvidia-ml-py` or the NVIDIA driver library is unavailable.

amd_gpus ¶

amd_gpus() -> list[GpuInfo]

Probe visible AMD GPUs through amdsmi, AMD's system management library.

Returns:

Type	Description
`list[GpuInfo]`	The visible AMD GPUs in processor-handle order, or an empty list when none are
`list[GpuInfo]`	present. Memory fields are byte counts and `architecture` is the gfx target.

Raises:

Type	Description
`ImportError`	If `amdsmi` or the AMD SMI driver library is unavailable.

intel_gpus ¶

intel_gpus() -> list[GpuInfo]

Probe visible Intel GPUs through pyzes, the Level Zero Sysman bindings.

Returns:

Type	Description
`list[GpuInfo]`	The visible Intel GPUs in Sysman device order, or an empty list when none are
`list[GpuInfo]`	present. Memory fields are byte counts, or `None` when Sysman reports no
`list[GpuInfo]`	memory modules. `compute_capability` and `architecture` are `None`.

Raises:

Type	Description
`ImportError`	If `pyzes` or the Level Zero driver library is unavailable.
`RuntimeError`	If a Level Zero Sysman query fails.

tpus ¶

tpus() -> list[TpuInfo]

Probe local TPU chips through tpu-info.

Detection is a local PCI scan; it needs neither libtpu nor a running runtime.

Returns:

Type	Description
`list[TpuInfo]`	A single-entry list describing the local TPU chip generation, its per-chip
`list[TpuInfo]`	specs, and count, or an empty list when no local TPU chips are present.

Raises:

Type	Description
`ImportError`	If `tpu-info` is unavailable; install `tinax[tpu]`.

GpuInfo `dataclass` ¶

GpuInfo(
    backend: GpuBackend,
    index: int,
    name: str,
    total_bytes: float | None,
    free_bytes: float | None,
    compute_capability: str | None = None,
    architecture: str | None = None,
)

A detected GPU; memory and capability fields are None when the vendor cannot report them.

compute_capability is NVIDIA's SM version and architecture is AMD's gfx target.

TpuInfo `dataclass` ¶

TpuInfo(
    generation: str,
    count: int,
    devices_per_chip: int,
    hbm_gib: int,
)

A detected local TPU chip generation, its per-chip specs, and how many are present.

generation is the lowercase chip name (for example "v5e"); devices_per_chip and hbm_gib are the fixed specs that tpu-info records for that generation.

Device¶

set_visible_cuda ¶

configure_single_chip ¶

configure_jax ¶

device_info ¶

DeviceInfo dataclass ¶

nvidia_gpus ¶

amd_gpus ¶

intel_gpus ¶

tpus ¶

GpuInfo dataclass ¶

TpuInfo dataclass ¶

DeviceInfo `dataclass` ¶

GpuInfo `dataclass` ¶

TpuInfo `dataclass` ¶