Skip to content

Compute hidden state updates in StatefulOnnxLabelScorer in batches#181

Draft
SimBe195 wants to merge 4 commits intomasterfrom
batched_hidden_state_updates
Draft

Compute hidden state updates in StatefulOnnxLabelScorer in batches#181
SimBe195 wants to merge 4 commits intomasterfrom
batched_hidden_state_updates

Conversation

@SimBe195
Copy link
Collaborator

@SimBe195 SimBe195 commented Feb 9, 2026

The most expensive part in search with an AED or LSTM LM is the hidden-state update which is currently done one-by-one for each hypothesis individually. This PR changes the logic to collect the unique hidden states first and then forward the state updater in a batched manner. Depending on how the state updater exports are structured, this may break older ONNX models since now the state updater needs to accept and return batched tensors (with batch axis 0).

So far not tested yet.

@larissakl
Copy link
Contributor

Is it maybe possible to offer both this batched version and a non-batched version if the state updater does not have a batch dimension?

@SimBe195
Copy link
Collaborator Author

Is it maybe possible to offer both this batched version and a non-batched version if the state updater does not have a batch dimension?

In principle this is possible but the only way to check for a batch dimension is looking up whether dim 0 is dynamic which might be unreliable and lead to edge cases.

@larissakl
Copy link
Contributor

Is it maybe possible to offer both this batched version and a non-batched version if the state updater does not have a batch dimension?

In principle this is possible but the only way to check for a batch dimension is looking up whether dim 0 is dynamic which might be unreliable and lead to edge cases.

Ideally, I would like to have a check if a "real" batch dimension exists and if yes use the batched version and if not automatically fall back to the unbatched version (maybe with a warning). But I see that this is not straightforward to realize...

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.

2 participants