# How to Debug JIT Compilation Errors and Profile JAX Training Pipelines in EvoRL

> Debug JIT errors and profile JAX training in EvoRL easily. Disable JIT, inspect compiled code with jax.debug.print, and analyze performance using jax.profiler and TensorBoard.

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: how-to-guide
- Published: 2026-03-01

---

**To debug JIT compilation errors and profile JAX training pipelines in EvoRL, disable JIT via the `enable_jit` config flag for step-through debugging, use `jax.debug.print` and `tree_has_nan` to inspect compiled code, and leverage `jax.profiler` with TensorBoard for performance analysis.**

EvoRL is an open-source evolutionary reinforcement learning framework built on JAX. Because EvoRL wraps its training loops, rollouts, and replay-buffer updates in `jax.jit` (or `jax.pmap` for multi-GPU) to achieve XLA-compiled speedups, understanding how to debug JIT compilation errors and profile JAX training pipelines in EvoRL is essential for stable experimentation.

## Common JIT Compilation Errors in EvoRL

### Static Argument Mismatches

The most frequent compilation failure occurs when arguments marked as static change value or type between calls. In [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py), the `enable_jit` property JIT-compiles `step` and `evaluate` with `static_argnums=(0,)`, ensuring the workflow instance itself remains static.

When a static argument becomes dynamic—such as passing a changing configuration dictionary—JAX raises a hashability error. Use `jax.make_jaxpr` on the failing function to inspect which inputs JAX treats as static versus dynamic.

### Shape Changes and Recompilation

If tensor shapes vary across training steps, JAX recompiles the function on every call, destroying performance. EvoRL encourages immutable shapes through its type system, but dynamic batch sizes or variable-length episodes can trigger this issue.

Insert `jax.debug.print` statements (as seen in commented lines within [`evorl/algorithms/offpolicy_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/offpolicy_utils.py)) to trace shape changes. Additionally, use `tree_has_nan` from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to detect NaN values that silently force recompilation.

### PyTree Structural Changes

EvoRL state objects inherit from `evorl/types.PyTreeData`, guaranteeing stable tree structures. Errors stating "Pytree structures differ" typically occur when dataclass fields marked `static: False` change type or when custom objects lack proper JAX registration.

Verify dataclass field definitions in [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py). Fields intended to remain static should be explicitly marked, while dynamic fields must maintain consistent shapes across iterations.

## Debugging Tools and Techniques

### Disabling JIT Compilation

For step-by-step debugging, disable JIT entirely via the configuration system. In [`configs/config.yaml`](https://github.com/emi-group/evorl/blob/main/configs/config.yaml), set `enable_jit: false`. The [`scripts/train.py`](https://github.com/emi-group/evorl/blob/main/scripts/train.py) entry point reads this flag and constructs the workflow accordingly, allowing standard Python debugging with `pdb` or IDE breakpoints.

```python

# In your config or script

config.enable_jit = False  # Disables JIT for debugging

workflow = RLWorkflow.build_from_config(config)

```

### Inspecting Static vs Dynamic Arguments

The `jit_method` decorator in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) wraps `jax.jit` with convenient static argument handling. To debug static argument issues:

```python
from evorl.utils.jax_utils import jit_method
import jax

@jit_method(static_argnums=(0,))
def my_step(config, state, action):
    return state.update(action)

# Inspect what JAX sees

jaxpr = jax.make_jaxpr(my_step)(config, state, action)
print(jaxpr)

```

### Detecting NaNs and Debug Printing

When values become NaN inside JIT-compiled code, standard Python `print` statements fail. Use `jax.debug.print` for host-side logging:

```python
import jax
from evorl.utils.jax_utils import tree_has_nan

@jax.jit
def train_step(state, batch):
    loss, grads = compute_loss(state, batch)
    
    # Debug prints (visible in host console)

    jax.debug.print("Loss value: {}", loss)
    jax.debug.print("NaN in grads: {}", tree_has_nan(grads))
    
    return state.apply_gradients(grads)

```

## Profiling JAX Training Pipelines

### Using jax.profiler with TensorBoard

EvoRL training scripts support JAX's native profiler for identifying compilation bottlenecks and device utilization issues. Wrap your training loop with `jax.profiler.start_trace` and `jax.profiler.stop_trace`:

```python
import jax
from evorl.workflows import RLWorkflow

def profile_training():
    # Initialize workflow

    workflow = RLWorkflow.build_from_config(config, enable_jit=True)
    state = workflow.init(jax.random.PRNGKey(0))
    
    # Start profiling trace

    jax.profiler.start_trace("/tmp/evorl_profile")
    
    # Run training loop

    state = workflow.learn(state)
    
    # Stop and save trace

    jax.profiler.stop_trace()
    print("Profile saved to /tmp/evorl_profile")

# Visualize with:

# tensorboard --logdir /tmp/evorl_profile

```

### Step-wise Performance Analysis

For granular timing without full profiling overhead, use `timeit` or `jax.experimental.host_callback` to measure individual step durations:

```python
import time
import jax

@jax.jit
def step_fn(state, action):
    return state.update(action)

# Warmup

_ = step_fn(state, action)

# Time 1000 steps

start = time.time()
for _ in range(1000):
    state = step_fn(state, action)
state.block_until_ready()  # Ensure async ops complete

elapsed = time.time() - start
print(f"Step time: {elapsed/1000*1000:.2f} ms")

```

## Configuration and Environment Setup

### GPU Memory and Deterministic Operations

Before launching training, configure JAX's XLA backend to prevent out-of-memory errors and ensure reproducibility. EvoRL provides utilities in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py):

```python
from evorl.utils.jax_utils import disable_gpu_preallocation, enable_deterministic_mode

# Disable GPU memory pre-allocation (helps with OOM errors)

disable_gpu_preallocation()

# Force deterministic GPU operations (slower but reproducible)

enable_deterministic_mode()

# Now safe to initialize workflow

workflow = RLWorkflow.build_from_config(config)

```

These functions set the `XLA_PYTHON_CLIENT_PREALLOCATE` and `XLA_FLAGS` environment variables automatically.

## Summary

- **Disable JIT** for step-through debugging by setting `enable_jit: false` in [`configs/config.yaml`](https://github.com/emi-group/evorl/blob/main/configs/config.yaml) or passing `enable_jit=False` to `RLWorkflow.build_from_config`.
- **Inspect static arguments** using `jax.make_jaxpr` and the `jit_method` decorator from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to resolve compilation mismatches.
- **Debug inside JIT** with `jax.debug.print` and `tree_has_nan` to catch NaNs and shape changes without breaking the compiled graph.
- **Profile performance** using `jax.profiler.start_trace` and TensorBoard to visualize XLA compilation times and device utilization.
- **Configure environment** with `disable_gpu_preallocation` and `enable_deterministic_mode` to prevent OOM errors and ensure reproducible results.

## Frequently Asked Questions

### How do I completely disable JIT compilation in EvoRL for debugging?

Set `enable_jit: false` in your [`configs/config.yaml`](https://github.com/emi-group/evorl/blob/main/configs/config.yaml) file or pass `enable_jit=False` when calling `RLWorkflow.build_from_config()`. This forces the workflow to execute pure Python, allowing you to use standard debuggers like `pdb` or IDE breakpoints inside [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py).

### Why does my EvoRL training loop recompile on every iteration?

Recompilation occurs when input shapes or PyTree structures change between calls, violating JAX's static shape requirement. Check for dynamic batch sizes or variable-length sequences. Use `jax.debug.print` to trace shapes and `tree_has_nan` from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to detect NaNs that trigger recompilation. Ensure all dataclass fields in [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py) maintain consistent static/dynamic markings.

### How can I detect NaN values inside a JIT-compiled function in EvoRL?

Use `jax.debug.print` to log values during execution, as standard Python print statements are unavailable inside JIT-compiled code. Import `tree_has_nan` from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to scan entire PyTrees for NaN values. Example usage includes printing loss values and gradient NaN status inside the training step function defined in [`evorl/algorithms/offpolicy_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/offpolicy_utils.py).

### What is the best way to profile GPU utilization in EvoRL training?

Wrap your training loop with `jax.profiler.start_trace("/tmp/profile")` before and `jax.profiler.stop_trace()` after execution. This captures XLA compilation times, kernel execution durations, and memory usage. Visualize results by running `tensorboard --logdir /tmp/profile`. For quick step-wise timing without full profiling overhead, use Python's `timeit` module with `block_until_ready()` to measure individual JIT-compiled step durations.