Allow TFTExplainer to explain more series than the model batch size
Hi @dennisbader,
I am opening this separately as requested in #2489.
Is your feature request related to a current problem? Please describe.
TFTExplainer.explain() currently raises a ValueError when a sequence contains more series than model.batch_size. The underlying prediction can process the sequence across several batches, but the TFT module retains its attention and variable-selection data only for the latest batch. Users with larger inputs must therefore split the target series and matching covariates themselves, call explain() for each chunk, and combine the results manually.
model.batch_size is the configured prediction batch size rather than an architectural limit on how many series prediction can process. This request is not for dataset-level global aggregation. The expected output remains one explanation result per input series.
Describe proposed solution
Allow TFTExplainer.explain() to process these inputs internally in chunks no larger than model.batch_size. Target series, past covariates, and future covariates would use the same chunk boundaries. After predicting each chunk, the explainer would collect its cached outputs before the next prediction overwrites them, then return all results in the original input order.
The existing result structure would be preserved for all six outputs.
- Attention
- Encoder importance
- Decoder importance
- Static covariates importance
- Encoder importance over time
- Decoder importance over time
Feature columns would remain consistently ordered across chunks. Models without static covariates would continue returning an empty (0, 0) static-importance DataFrame for each series. Single-series calls and inputs that already fit within one batch would remain unchanged.
Describe potential alternatives
Users can continue calling explain() once per chunk, or increase model.batch_size to cover every series when memory allows. The proposed implementation would make ceil(number_of_series / model.batch_size) prediction calls. This repeats fixed prediction setup for each chunk, but keeps batch memory bounded and avoids a more invasive change to collect every batch cache during a single prediction.
Additional context
This follows the discussion in #2489 and the request to track internal batching separately. #2955 documents why the cached data only covers the final prediction batch, and #2957 added the current guard. #3170 is now merged and adds the encoder and decoder importances over time that must also be collected for every chunk.
I have a local implementation against current master at f1e421c7. With five distinct series and batch_size=2, the combined result from chunks of 2, 2, and 1 matched five singleton explain() calls across all six outputs. The checks cover TimeSeries values, time indices, and components, as well as DataFrame values, row indices, and consistently ordered feature columns. A separate no-static-covariate test confirms that every series retains an empty (0, 0) DataFrame.
The 54 TFT explainer tests and the full 168-test explainability suite pass. pre-commit also passes.
I can open a PR with this implementation if this approach looks good.
Source: unit8co/darts