Skip to content

feat(inference): add average_samples flag to auto_regressive_inference#321

Open
moscowmule2240 wants to merge 1 commit into
shiyu-coder:masterfrom
moscowmule2240:feat/return-individual-samples
Open

feat(inference): add average_samples flag to auto_regressive_inference#321
moscowmule2240 wants to merge 1 commit into
shiyu-coder:masterfrom
moscowmule2240:feat/return-individual-samples

Conversation

@moscowmule2240
Copy link
Copy Markdown

What

Add an average_samples flag (default True) to auto_regressive_inference.

  • True (default): unchanged — averages over the sample dimension, returns (batch, total_seq, n_feat).
  • False: returns per-sample trajectories (batch, sample_count, total_seq, n_feat), exposing the full Monte Carlo distribution.

Why

Obtaining individual sampled paths (probabilistic forecasting / uncertainty / quantile bands) currently requires re-implementing the autoregressive loop, since the function always collapses the sample dimension via np.mean. This is a one-line, fully backward-compatible change.

Tests

tests/test_average_samples.py (fake tokenizer/model, no weight download): shape behavior for both flag values + averaged == mean over samples.

By default (average_samples=True) behavior is unchanged: predictions are
averaged over the sample dimension. With average_samples=False the per-sample
paths are returned with shape (batch, sample_count, total_seq, n_feat), letting
callers obtain the full Monte Carlo distribution (probabilistic forecasting /
uncertainty) without re-implementing the autoregressive loop. Backward compatible.

Adds a docstring for the flag and tests/test_average_samples.py (fake model, no
weight download) covering the shape behavior and mean-over-samples consistency.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant