flashquad.utils.backend

Backend resolution and dtype validation.

Module Contents

Functions

_resolve_backend_namespace	Resolve a backend string to its array namespace.
_default_dtype	Return the namespace-native default floating dtype.
_validate_dtype	Raise TypeError when dtype is incompatible with backend_name.
_resolve_device	Resolve compute device for the given backend.
_to_device	Move tensor to device when the backend supports it.

Data

_BACKEND_MODULES
_TORCH_BACKENDS
_NUMPY_COMPAT_BACKENDS

API

flashquad.utils.backend._BACKEND_MODULES

None

flashquad.utils.backend._TORCH_BACKENDS

‘frozenset(…)’

flashquad.utils.backend._NUMPY_COMPAT_BACKENDS

‘frozenset(…)’

flashquad.utils.backend._resolve_backend_namespace(backend)

Resolve a backend string to its array namespace.

flashquad.utils.backend._default_dtype(xp)

Return the namespace-native default floating dtype.

flashquad.utils.backend._validate_dtype(backend_name, dtype)

Raise TypeError when dtype is incompatible with backend_name.

flashquad.utils.backend._resolve_device(backend_name, device)

Resolve compute device for the given backend.

For 'torch' and 'jax', auto-selects GPU when available if device is None. Other backends return None (cupy is inherently GPU; numpy is CPU-only).

flashquad.utils.backend._to_device(tensor, device)

Move tensor to device when the backend supports it.