Custom Kernel Framework

Custom Kernel

XLACustomKernel(name[, doc])	Creates and manages a custom JAX primitive for XLA custom calls.
KernelEntry(backend, platform, kernel_generator)	A registered kernel implementation for a specific backend and platform.

CPU Kernel via Numba

numba_kernel(kernel, outs, *[, vmap_method, ...])

Create a JAX-callable function from a Numba CPU kernel.

GPU Kernel via Numba CUDA

numba_cuda_kernel(kernel, outs, *[, grid, ...])	Create a JAX-callable function from a single Numba CUDA kernel.
numba_cuda_callable(func, outs, *[, ...])	Create a JAX-callable from a Python function that launches Numba CUDA kernels.

GPU Kernel via CUDA Source

load_cuda_inline(name, cuda_sources[, ...])	Compile inline CUDA source and load the resulting module.
load_cuda_file(filepath[, functions, name])	Compile a single .cu file and load the resulting module.
load_cuda_dir(directory[, functions, name, ...])	Compile all CUDA files in a directory and load the resulting module.

CPU Kernel via C++ Source

load_cpp_inline(name, cpp_sources[, ...])	Compile inline C++ source for CPU (or CUDA) and load the module.
load_cpp_file(filepath[, functions, name])	Compile a single .cpp / .cc file for CPU and load the module.
normalize_tokens(tokens)	Normalise compatible tokens to canonical BE form.

Compiler Backends

CompilerBackend()	Abstract base class for BE compiler backends.
CPPBackend(toolchain)	Compile C++ sources with g++ / clang++.
CUDABackend(toolchain)	Compile CUDA sources directly with nvcc.
HIPBackend()	Compile HIP sources for AMD GPUs (stub — not yet implemented).

Runtime

CompiledModule(so_path, function_names)

A compiled module loaded from a shared library.

register_ffi_target(target_name, module, ...)	Register a compiled function as a JAX FFI target.
list_registered_targets()	Return a sorted list of all registered FFI target names.

Primitive Registry

get_registry()	Return a copy of the full primitive registry.
get_all_primitive_names()	Return a sorted list of all registered primitive names.
get_primitives_by_tags(tags)	Return primitives that have all the specified tags.

Cache Utilities

clear_cache([name])	Clear the compilation cache.
set_cache_dir(path)	Set the compilation cache directory.
get_cache_dir()	Return the current compilation cache directory path.

Diagnostics

print_diagnostics()

Print a summary of the brainevent compilation environment.

Exceptions

Kernel construction, compilation and dispatch raise the exceptions documented in Error Classes.