spectralbrain.backends.gpu#

GPU compute backends — CuPy, JAX, NumPyro, and VRAM management.

This module provides:

CupyBackend — drop-in GPU replacement for NumpyBackend using CuPy (cupy-cuda13x).
JaxBackend — GPU backend with jit and vmap for batch-subject spectral descriptor computation.
NumPyroSampler — GPU-accelerated Bayesian MCMC via JAX.
VRAM management — monitoring, cache clearing, defragmentation, garbage collection, and a memory-guarded context manager.

All dependencies are lazy-imported. If CuPy or JAX is missing, the module still imports successfully — only instantiation of the backends raises ImportError.

Functions

`get_gpu_backend`([name])	Factory for GPU compute backends.
`get_gpu_bayesian_sampler`([backend])	Factory for GPU Bayesian samplers.
`vram_clear`()	Clear CUDA memory caches across all known GPU frameworks.
`vram_defrag`([device_id])	Attempt CUDA memory defragmentation.
`vram_gc`([device_id])	Full GPU garbage collection: Python GC + VRAM clear + defrag.
`vram_guard`([min_free_gb, device_id, ...])	Context manager for VRAM-safe GPU operations.
`vram_monitor`([device_id, label])	Log current VRAM usage (one-shot, for debugging).
`vram_status`([device_id])	Query current VRAM usage.

Classes

`BlackjaxSampler`([num_warmup, num_samples, ...])	GPU-accelerated Bayesian NUTS via BlackJAX.
`CupyBackend`([device_id])	GPU compute backend using CuPy.
`JaxBackend`([device])	GPU backend using JAX with `jit` and `vmap`.
`NumPyroSampler`([num_warmup, num_samples, ...])	GPU-accelerated Bayesian MCMC using NumPyro + JAX.
`TorchBackend`([device])	GPU compute backend using PyTorch.
`VRAMInfo`(device_name, device_id, total_gb, ...)	Snapshot of GPU VRAM usage.

class spectralbrain.backends.gpu.BlackjaxSampler(num_warmup=1000, num_samples=2000, num_chains=4, seed=42)[source]#

Bases: object

GPU-accelerated Bayesian NUTS via BlackJAX.

BlackJAX is a low-level sampler that operates on a log-density function rather than a model object, which makes it composable and fast under JAX jit/vmap on the GPU. This wrapper runs the standard window-adaptation → NUTS pipeline and (optionally) vectorises independent chains with jax.vmap().

Parameters:

num_warmup (int) – Window-adaptation (tuning) steps.
num_samples (int) – Posterior draws per chain.
num_chains (int) – Independent chains, run in parallel via vmap.
seed (int) – PRNG seed.

Examples

>>> import jax.numpy as jnp
>>> def logdensity(theta):
...     # standard-normal target
...     return -0.5 * jnp.sum(theta ** 2)
>>> sampler = BlackjaxSampler(num_warmup=500, num_samples=1000)
>>> samples = sampler.sample(logdensity, initial_position=jnp.zeros(3))
>>> samples.shape  # (num_samples, 3)

sample(logdensity_fn, initial_position, *, rng_key=None)[source]#

Run NUTS on a log-density function.

Parameters:

logdensity_fn (callable) – Maps a parameter pytree to a scalar log-density (unnormalised log-posterior). Must be JAX-traceable.
initial_position (pytree) – Starting position for a single chain. For num_chains > 1 it is broadcast across chains.
rng_key (jax.Array, optional) – PRNG key. Defaults to jax.random.PRNGKey(self.seed).

Returns:

pytree – Posterior draws. For a single chain each leaf has shape (num_samples, *param_shape); for multiple chains (num_chains, num_samples, *param_shape).

Return type: