Skip to content

Batch inference

Fit many independent models in one call. rustmc runs their chains through a shared Rayon thread pool, and offers two batch shapes:

  • CompiledModel.sample_batch() reuses one immutable graph structure across validated datasets that share a schema. Use this when the model is the same and only the data changes -- one demand model per SKU, one calibration per instrument.
  • rustmc.batch_sample() accepts a different model structure per entry, so each entry owns its own graph and dataset. Use it only when the structures really do differ.

This example uses 100 SKUs so it stays practical to run locally. It is an API example, not a throughput claim. Whether rustmc, ARIMA or Prophet is faster depends on the model, configuration, data and hardware, and their default uncertainty outputs are not directly comparable. benchmarks/comparisons/batch_many_series.py has an exploratory rustmc/PyMC+nutpie comparison that reports divergences, R-hat and ESS/s next to wall time; benchmarks/README.md says what may be claimed from a measurement.

Run it with python examples/batch_inference.py from the repository root. The blocks below are cells of that one file and share its state, so they assume the blocks above them have run.

import numpy as np

import rustmc as rmc

Simulate 100 weekly SKU series

np.random.seed(0)
N_MODELS = 100  # benchmark larger runs on your own model and hardware
T = 52  # weeks per SKU

true_intercepts = np.random.normal(100, 20, N_MODELS)
true_trends = np.random.normal(0.5, 0.2, N_MODELS)
noise_std = 5.0

# One time axis, in weeks, shared by every series. The data is generated against
# this same axis, so a fitted `trend` is directly comparable to `true_trends`.
t = np.arange(T, dtype=np.float64)

datasets = [
    {
        "t": t,
        "y": true_intercepts[i] + true_trends[i] * t + np.random.normal(0, noise_std, T),
    }
    for i in range(N_MODELS)
]

print(f"{N_MODELS} series, {T} weeks each")
print(f"true intercept: mean {true_intercepts.mean():.2f}, sd {true_intercepts.std():.2f}")
print(f"true trend:     mean {true_trends.mean():.3f}, sd {true_trends.std():.3f}  (units per week)")

100 series, 52 weeks each
true intercept: mean 101.20, sd 20.16
true trend:     mean 0.516, sd 0.207  (units per week)

Shared structure: CompiledModel.sample_batch()

Every SKU here has the same model, so the graph is built and compiled once and each dataset is bound to it. errors="collect" attempts every cell and returns the failures in batch.errors keyed by dataset ID, instead of losing the whole batch to one bad series.

t starts at zero, so the intercept is the level in week 0 and is correlated with the trend. The cell after the fit measures that correlation rather than asserting it. Centering time -- t - t.mean() -- removes it, at the cost of an intercept that means "level at mid-year".

Fit all 100 SKUs

builder = rmc.ModelBuilder()
intercept = builder.normal_prior("intercept", mu=0.0, sigma=200.0)
trend = builder.normal_prior("trend", mu=0.0, sigma=20.0)
builder.normal_likelihood("obs", mu_expr=intercept + trend * "t", sigma=noise_std, observed_key="y")
compiled = builder.compile()

batch = compiled.sample_batch(
    datasets,
    ids=[f"sku-{i:03d}" for i in range(N_MODELS)],
    chains=1,
    draws=500,
    warmup=300,
    seed=42,
    errors="collect",
    show_progress=False,
)

print(f"fitted {len(batch)} datasets, {len(batch.errors)} failed")
for failed_id, message in batch.errors.items():
    print(f"  {failed_id}: {message}")

# `batch.errors` is empty here, but reading a failed cell raises, so everything
# below works from the IDs that succeeded rather than from range(N_MODELS).
ok = [i for i, name in enumerate(batch.ids) if name not in batch.errors]

fitted 100 datasets, 0 failed

Compare the first five to the values that generated them

print(f"{'SKU':<9} {'intercept':>19} {'true':>9} {'trend':>18} {'true':>9}")
for i in ok[:5]:
    fit = batch[i]
    mean, std = fit.mean(), fit.std()
    print(
        f"{batch.ids[i]:<9} "
        f"{mean['intercept']:9.2f} +/- {std['intercept']:5.2f} {true_intercepts[i]:9.2f} "
        f"{mean['trend']:9.3f} +/- {std['trend']:5.3f} {true_trends[i]:9.3f}"
    )

SKU                 intercept      true              trend      true
sku-000      135.77 +/-  1.48    135.28     0.833 +/- 0.050     0.877
sku-001      106.28 +/-  1.18    108.00     0.298 +/- 0.042     0.230
sku-002      118.32 +/-  1.37    119.57     0.248 +/- 0.047     0.246
sku-003      143.51 +/-  1.46    144.82     0.714 +/- 0.048     0.694
sku-004      137.82 +/-  1.34    137.35     0.246 +/- 0.045     0.265

Batch-wide recovery and diagnostics

fitted_intercepts = np.array([batch[i].mean()["intercept"] for i in ok])
fitted_trends = np.array([batch[i].mean()["trend"] for i in ok])
divergences = np.array([batch[i].divergences for i in ok])
intercept_error = fitted_intercepts - true_intercepts[ok]
trend_error = fitted_trends - true_trends[ok]

print(f"intercept error: mean {intercept_error.mean():+.3f}, "
      f"rmse {np.sqrt(np.mean(intercept_error**2)):.3f}")
print(f"trend error:     mean {trend_error.mean():+.4f}, "
      f"rmse {np.sqrt(np.mean(trend_error**2)):.4f}")
print(f"divergences:     {divergences.sum()} across {len(ok)} fits "
      f"({int((divergences > 0).sum())} fits affected)")

intercept error: mean -0.050, rmse 1.243
trend error:     mean -0.0020, rmse 0.0386
divergences:     0 across 100 fits (0 fits affected)

What a BatchResult carries

r = batch[ok[0]]
print("mean()                ", {k: round(v, 3) for k, v in r.mean().items()})
print("std()                 ", {k: round(v, 3) for k, v in r.std().items()})
print("get_samples()         ", {k: v.shape for k, v in r.get_samples().items()})
print("get_samples_2d()      ", {k: v.shape for k, v in r.get_samples_2d().items()})
print("accept_rate           ", round(r.accept_rate, 3))
print("accept_rates          ", [round(a, 3) for a in r.accept_rates])
print("divergences           ", r.divergences)
print("divergences_per_chain ", r.divergences_per_chain)

# The intercept/trend correlation this parameterization implies, measured.
draws = r.get_samples()
print("corr(intercept, trend)", round(float(np.corrcoef(draws["intercept"], draws["trend"])[0, 1]), 3))

mean()                 {'intercept': 135.767, 'trend': 0.833}
std()                  {'intercept': 1.475, 'trend': 0.05}
get_samples()          {'intercept': (500,), 'trend': (500,)}
get_samples_2d()       {'intercept': (1, 500), 'trend': (1, 500)}
accept_rate            0.93
accept_rates           [0.93]
divergences            0
divergences_per_chain  [0]
corr(intercept, trend) -0.885

chains=1 is the throughput-first setting, and it gives up R-hat, which needs more than one chain. Raise chains when per-model convergence evidence matters more than batch wall time.

Different structures: rmc.batch_sample()

# Each entry owns its own graph, so the entries need not share a schema. The third
# series below ran promotions in two four-week blocks and gets a term for them that
# the other two do not have. The blocks sit in the middle of the year rather than at
# the end, so the promotion indicator is not confounded with the trend.
promo = (((t >= 12) & (t < 16)) | ((t >= 34) & (t < 38))).astype(np.float64)
true_lift = 12.0
promo_y = true_intercepts[2] + true_trends[2] * t + true_lift * promo + np.random.normal(0, noise_std, T)

models = []
for data in [datasets[0], datasets[1], {"t": t, "promo": promo, "y": promo_y}]:
    entry = rmc.ModelBuilder()
    a = entry.normal_prior("intercept", mu=0.0, sigma=200.0)
    b = entry.normal_prior("trend", mu=0.0, sigma=20.0)
    mu_expr = a + b * "t"
    if "promo" in data:
        lift = entry.normal_prior("lift", mu=0.0, sigma=50.0)
        mu_expr = mu_expr + lift * "promo"
    entry.normal_likelihood("obs", mu_expr=mu_expr, sigma=noise_std, observed_key="y")
    models.append((entry.build(), data))

results = rmc.batch_sample(models, chains=1, draws=500, warmup=300, seed=42, show_progress=False)
for i, result in enumerate(results):
    mean, std = result.mean(), result.std()
    params = ", ".join(f"{k}={mean[k]:.3f} +/- {std[k]:.3f}" for k in sorted(mean))
    print(f"model {i}: {params}")
print(f"true lift on model 2: {true_lift:.1f}")

model 0: intercept=135.542 +/- 1.502, trend=0.841 +/- 0.050
model 1: intercept=106.152 +/- 1.319, trend=0.302 +/- 0.044
model 2: intercept=123.313 +/- 1.479, lift=10.622 +/- 1.955, trend=0.133 +/- 0.048
true lift on model 2: 12.0

Notes

  • Every model in a batch shares one draws and warmup count.
  • The thread pool is shared across chains and models. More cores may cut wall time, but scaling depends on model size, batch size, memory bandwidth and scheduling. Measure it on the workload you care about rather than assuming it is proportional.
  • sampler="hmc" is available in batch mode as a fixed-step fallback.
  • Chunked batch dispatch currently retains inputs and fits, so peak memory grows with the batch. See the roadmap for bounded streaming.