batched SSM compatibility#685

Draft

Dekermanjian wants to merge 10 commits into

pymc-devs:mainfrom

Dekermanjian:ssm_batch_dims

Commits on May 20, 2026

rough implementation of batched SSM compatability
Dekermanjian
committed

Commits on May 21, 2026

added a parameter to explicitly define the batch size, updated VARMAX, SARIMAX, ETS so that positional arguments don't clash with the new parameter, updated error in control flow for data registry
Dekermanjian
committed

Commits on May 26, 2026

refactored vectorization methods, pulled out repeated helpers into batch_tools utilities file, use slicing logic instead of branching for forecast indexing, make batch_size a tuple for destructurin…
Dekermanjian
committed

Commits on May 29, 2026

always call vectorize_graph, build_statespace is entrypoint for batch_sizes, added a utility to infer the batch dimension
Dekermanjian
committed

Commits on Jun 7, 2026

1. fixed a bug with sample_filter_outputs() when vector_graph doesn't vectorize certain matrices
Dekermanjian
committed

Commits on Jun 11, 2026

removed mock_sample from batch_ssm tests
Dekermanjian
committed