brainevent.config.set_backend

Contents

brainevent.config.set_backend#

brainevent.config.set_backend(platform, backend)[source]#

Set the global default backend for a platform across all primitives.

After this call, every primitive that has a kernel registered for backend on platform will use it by default, unless overridden by an explicit backend= keyword argument at call time.

Parameters:

  • platform (str) – The platform name (e.g., 'cpu', 'gpu', 'tpu').

  • backend (str | None) – The backend name (e.g., 'warp', 'pallas', 'numba'). Pass None to clear the global default for this platform, reverting to per-primitive defaults.

Raises:

ValueError – If backend is an empty string.

See also

get_backend

Query the current global backend for a platform.

clear_backends

Clear all global backend defaults.

Notes

Backend selection is resolved inside each primitive’s MLIR lowering function, and JAX caches lowered/compiled executables (both the eager dispatch cache and the jit compilation cache) keyed on primitive + abstract values, not on the global backend setting. If this call actually changes the effective value for platform, it therefore calls jax.clear_caches() so that already-compiled call sites pick up the new backend on their next invocation (at the cost of forcing recompilation everywhere on next use). Calling set_backend with the value it already has is a no-op and does not clear caches.

This makes set_backend a setup-time, single-threaded control: it is intended to be called before the hot loop / before other threads start compiling, not toggled concurrently from multiple threads while other work is in flight (jax.clear_caches() is process-global). For per-call, thread-safe backend selection, pass backend=<name> directly to the primitive call instead – it is a bind parameter and therefore part of the cache key, so it composes safely with concurrent dispatch.

Examples

>>> import brainevent
>>> brainevent.set_backend('gpu', 'warp')
>>> brainevent.get_backend('gpu')
'warp'
>>> brainevent.set_backend('gpu', None)  # clear
>>> brainevent.get_backend('gpu') is None
True
Contents