transformers - 💡(How to fix) Fix docs: clarify output_hidden_states behavior for Koopman/spectral analysis use cases

Summary

When using output_hidden_states=True in AutoModel.forward(), the documentation does not clearly explain:

1. Which layers are included: Whether layer 0 is the embedding layer or the first transformer block varies by architecture.
2. Shape conventions: (batch, seq_len, hidden_dim) per layer, but the tuple indexing (offset-by-one vs direct) is architecture-dependent.
3. Recommended pattern for trajectory extraction: For dynamical systems analysis (e.g., Koopman/DMD spectral auditing), users need to stack hidden states across layers into a trajectory matrix. There is no canonical example for this.

Motivation

This was encountered while building koopgauge, a Koopman/DMD spectral audit tool for foundation models. The tool extracts hidden state trajectories and applies Dynamic Mode Decomposition to detect representation collapse and periodicity.

The current docs require users to empirically determine layer indexing conventions per model, which is fragile and leads to off-by-one bugs.

Proposed improvement

Add a section to the output_hidden_states documentation (or PreTrainedModel docstring) clarifying:

# output_hidden_states=True returns a tuple of length num_hidden_layers + 1
# Index 0: embedding output (before transformer blocks)
# Index i (i >= 1): output of transformer block i-1
outputs = model(**inputs, output_hidden_states=True)
all_hidden = outputs.hidden_states  # tuple[Tensor], len = num_layers + 1
# Stack into trajectory: shape (num_layers+1, seq_len, hidden_dim)
trajectory = torch.stack(all_hidden, dim=0)

This would improve usability for interpretability and analysis tools.

References

• koopgauge — Koopman/DMD spectral audit for sequence models
• RKSP: Kim et al. (2026), arXiv 2602.22988 — Runtime Koopman Spectral Profiling of LLMs