Skip to content

pixano_inference.frameworks.jax

JAX framework adapter.

Provides framework-neutral device/dtype resolution and numpy interchange for custom models written in JAX. JAX is imported lazily; importing this module does not require it.

JaxAdapter

:class:~pixano_inference.frameworks.base.FrameworkAdapter for JAX.

from_numpy(array, device=None, dtype=None)

Convert a numpy array to a JAX array on device with dtype.

Source code in pixano_inference/frameworks/jax.py 
60
61
62
63
64
65
66
67
68
def from_numpy(self, array: np.ndarray, device: Any = None, dtype: Any = None) -> Any:
    """Convert a numpy array to a JAX array on *device* with *dtype*."""
    import jax
    import jax.numpy as jnp

    result = jnp.asarray(array, dtype=dtype)
    if device is not None:
        result = jax.device_put(result, device)
    return result

is_available()

Whether jax is importable.

Source code in pixano_inference/frameworks/jax.py 
29
30
31
def is_available(self) -> bool:
    """Whether jax is importable."""
    return is_jax_installed()

resolve_device(num_gpus)

Return a JAX device: a GPU/TPU device when requested and present, else CPU.

Source code in pixano_inference/frameworks/jax.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
def resolve_device(self, num_gpus: float) -> Any:
    """Return a JAX device: a GPU/TPU device when requested and present, else CPU."""
    import jax

    if num_gpus > 0:
        for platform in ("gpu", "tpu"):
            try:
                devices = jax.devices(platform)
            except RuntimeError:
                devices = []
            if devices:
                return devices[0]
    return jax.devices("cpu")[0]

resolve_dtype(dtype)

Resolve a jax.numpy dtype from a string (e.g. "float32", "bfloat16").

Source code in pixano_inference/frameworks/jax.py 
47
48
49
50
51
52
53
54
def resolve_dtype(self, dtype: str) -> Any:
    """Resolve a jax.numpy dtype from a string (e.g. ``"float32"``, ``"bfloat16"``)."""
    import jax.numpy as jnp

    try:
        return jnp.dtype(dtype)
    except TypeError as exc:
        raise ValueError(f"Unsupported jax dtype '{dtype}'.") from exc

to_numpy(array)

Convert a JAX array to a numpy array (host memory).

Source code in pixano_inference/frameworks/jax.py 
56
57
58
def to_numpy(self, array: Any) -> np.ndarray:
    """Convert a JAX array to a numpy array (host memory)."""
    return np.asarray(array)