Device¶

tinax.device owns backend-agnostic pre-JAX environment policy that the JAX runtime honors, a summary of the active JAX runtime, and optional accelerator detection gated behind the gpu and tpu 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 needs the matching extra (pip install tinax[gpu] or tinax[tpu]):

from tinax.device import gpus, tpus

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

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.

gpus ¶

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 if none are
`list[GpuInfo]`	present. Memory fields are byte counts.

Raises:

Type	Description
`ImportError`	If `nvidia-ml-py` is unavailable; install `tinax[gpu]`.

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(
    index: int,
    name: str,
    total_bytes: float,
    free_bytes: float,
    compute_capability: str,
)

A detected NVIDIA GPU; compute_capability is comparable via float().

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 ¶

gpus ¶

tpus ¶

GpuInfo dataclass ¶

TpuInfo dataclass ¶

DeviceInfo `dataclass` ¶

GpuInfo `dataclass` ¶

TpuInfo `dataclass` ¶