stormlog.jax.jax_env

JAX environment configuration for Stormlog.

Suppresses verbose JAX/XLA logging and configures the JAX runtime environment before any import jax call. Every module in the stormlog.jax package should call configure_jax_logging() at import time, before importing jax itself.

Functions

configure_jax_logging()

Suppress verbose JAX/XLA info-level logging.
stormlog.jax.jax_env.configure_jax_logging()[source]

Suppress verbose JAX/XLA info-level logging.

Idempotent — safe to call multiple times. Sets environment variables that JAX and XLA inspect on first import:

  • JAX_LOG_COMPILES → "0" (suppress JIT compilation logs)

  • TF_CPP_MIN_LOG_LEVEL → "2" (suppress TF C++ backend noise when JAX falls back to the TF XLA bridge)

Return type:

None