tud_lbm.config.jax_config

JAX configuration settings for TUD-LBM.

This module provides centralized JAX configuration options. Import and call configure_jax() at the start of your script to apply settings.

Functions

configure_jax(→ None)

Configure JAX settings.

Module Contents

tud_lbm.config.jax_config.configure_jax(enable_x64: bool | None = None, disable_jit: bool | None = None) → None[source]

Configure JAX settings.

Call this function at the start of your script to apply JAX configuration. If parameters are None, the module-level defaults are used.

Parameters:

  • enable_x64 – Enable 64-bit precision. Defaults to module-level ENABLE_X64.

  • disable_jit – Disable JIT compilation for debugging. Defaults to module-level DISABLE_JIT.