Is callbacksupported for numpyro NUTS sampler?

[AlfredoJF]

I want to add an early stopping callback but seems numpyro sampler is not supported only pymc

Does anybody if callback in the sample method is supported for numpyro? Any example?

Thanks!

[ricardoV94]

Not that I know, but that would be a numpyro question not pymc. You can ask them in https://forum.pyro.ai/c/numpyro/9

[Vital-Fernandez]

Hi @alfredoj (and @ricardoV94)! I wondered if you found a solution: I am using pymc through a streamlit interface and I can make use a callback to track the progress:

draws, tune, chains, cores = 1000, 2000, 8, 8

total_steps = draws * chains  # callback only fires on draws, not tune
chain_draws = {i: 0 for i in range(chains)}
progress_bar = st.progress(0, text="Sampling...")

def sampling_callback(trace, draw):
    if not draw.tuning:
        chain_draws[draw.chain] += 1
        completed = sum(chain_draws.values())
        progress = min(completed / total_steps, 1.0)
        progress_bar.progress(progress,
                              text=f"Chain {draw.chain + 1}/{chains} | draw {chain_draws[draw.chain]}/{draws}")

# <- Run model function here

progress_bar.progress(1.0, text="Sampling complete!")

However, it does not work for numpyro. Could there be another way to hint the user that it is progressing?