Validate generation subgraph shapes

[Akshay Sonawane (apsonawane)]

This pull request strengthens input validation and error handling for subgraph classes in the ONNX Runtime transformers codebase. It introduces additional checks for required input/output shapes and dimensions, ensuring that missing or malformed shapes are detected early with clear error messages. The changes also add corresponding unit tests to verify the new validation logic.

Validation and Error Handling Improvements:

• Added checks to ensure subgraph output names are not empty in subgraph_base.cc, and that required tensor shapes (such as logits and past state shapes) are not null before proceeding with parameter extraction or validation. This applies to GPT, T5, and Whisper subgraph classes (subgraph_base.cc, subgraph_gpt.cc, subgraph_t5_decoder.cc, subgraph_t5_encoder.cc, subgraph_whisper_decoder.cc, subgraph_whisper_encoder.cc). [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12]

• Improved validation of input dimensions, such as ensuring decoder subgraph input tensors have the expected number of dimensions for T5 and Whisper models (subgraph_t5_decoder.cc, subgraph_whisper_decoder.cc). [1] [2]

Unit Testing:

• Added new unit tests in beam_search_test.cc to verify that missing logits or past state shapes are correctly rejected, and that error messages are descriptive. This includes helper code to load test models and manipulate subgraph outputs for testing. [1] [2] [3]

Bug Fixes:

• Fixed an error message in subgraph_whisper_encoder.cc to reference the correct input index for decoder_input_ids.

These changes improve robustness by ensuring that subgraph classes fail early and clearly when required inputs are missing or malformed, and they are now covered by dedicated unit tests.

[Copilot]

Pull request overview

This pull request hardens shape/dimension validation for transformer generation subgraphs (GPT/T5/Whisper) so malformed or missing tensor shapes are rejected early with clearer Status errors, and adds unit tests to ensure these failures remain covered.

Changes:

• Added new null/shape/dimension validation checks in multiple subgraph Validate paths and in Subgraph::GetParameters/Subgraph::Setup.
• Improved specific error messages (e.g., corrected Whisper encoder input index) and added explicit logits-shape null checks before parameter extraction.
• Added beam search unit tests that mutate model output shape metadata to validate the new error paths.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/contrib_ops/cpu/transformers/subgraph_base.cc	Adds early validation for required outputs and rejects null past-shape before extracting parameters.
onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc	Adds explicit rejection when logits output shape is missing.
onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc	Adds minimum-input and input_ids rank validation; rejects missing logits shape.
onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.cc	Rejects missing logits shape prior to parameter extraction.
onnxruntime/contrib_ops/cpu/transformers/subgraph_whisper_decoder.cc	Adds minimum-input and input_ids rank validation; rejects missing logits shape.
onnxruntime/contrib_ops/cpu/transformers/subgraph_whisper_encoder.cc	Fixes an input-index error message and rejects missing logits shape.
onnxruntime/test/contrib_ops/beam_search_test.cc	Adds unit tests ensuring missing logits/past shapes are rejected with descriptive errors.

Suppressed comments (1)

onnxruntime/contrib_ops/cpu/transformers/subgraph_whisper_encoder.cc:52

• WhisperEncoderSubgraph::Validate checks that input 0 is named "encoder_input_ids", but the subgraph documentation (and the later type-check message) indicate the input is "encoder_input_features" (float). As written, a correctly-named Whisper encoder subgraph would be rejected during validation.

Status WhisperEncoderSubgraph::Validate(const std::vector<const NodeArg*>& subgraph_inputs,
                                        const std::vector<const NodeArg*>& subgraph_outputs) {
  ORT_RETURN_IF(num_subgraph_inputs != 2, "expect 2 inputs, got:", num_subgraph_inputs);

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.