fix: bound the memory of long batched MCMC runs by clearing JAX's compilation caches

#!/bin/bash

# Copy the working tree of the laptop's repository (the refactor/samplers branch, committed or not) to the server

# clone used for tests and benchmarks (GitHub and GitLab clones fail on the server). Run outputs, checkpoints and

# node_modules of the configurator test are synced too, except the large posterior chains.

# node_modules of the configurator test are synced too, except the large posterior chains. What the server produces

# itself (run folders, results) is protected from deletion: an earlier version of this script, with --delete, wiped

# the outputs of a finished queue (2026-10-06).

# usage: sync_refactor_server.sh [server clone, default /DATA/ExoREC/pyaetna_bench/PyAETNA_refactor]

set -u

R=$(realpath "$(dirname "$0")/../..")

D=${1:-/DATA/ExoREC/pyaetna_bench/PyAETNA_refactor}

rsync -a --delete --exclude .git --exclude '__pycache__' --exclude '*.egg-info' --exclude 'paper/' \

rsync -a --delete --filter='P /paper_benchmarks/*/runs/**' --filter='P /paper_benchmarks/*/results/**' \

  --filter='P /paper_benchmarks/*/runs/*/**' --exclude .git --exclude '__pycache__' --exclude '*.egg-info' --exclude 'paper/' \

  --exclude 'juliet_out_*' --exclude '*.pckl' --exclude 'prompt.md' \

  "$R/" "ganymede-cuda13:$D/" 2>&1 | grep -v 'bind\|channel\|forward'

echo "synced $(git -C "$R" rev-parse --short HEAD) + working tree → $D"

"""MCMC sampling execution helpers and kernel factories."""

import gc

import time as ttt

from collections.abc import Callable

from typing import Any

from .context import SamplingContext

from .sampling_logging import _format_rate, format_time

CLEAR_CACHES_EVERY = (

    20  # batches between clearings of JAX's in-memory compilation caches

)

def run_batched(

    mcmc_obj: Any,

            mcmc_last_state=last_state,

            log_func=log_func,

        )

        # NumPyro compiles its sampling loop again at every run (a new closure each time); the executables pile up in

        # JAX's in-memory caches (about 60 MB per run for K2-140), which ran long ensemble runs out of memory.

        # Clearing the caches now and then bounds the memory; the next run reloads from the persistent cache.

        if batch_index % CLEAR_CACHES_EVERY == 0:

            jax.clear_caches()

            gc.collect()

        if (

            stop_when is not None

            and samples_so_far < target