Use dynamic shape storage in BeamSearch ExpandBuffer

[Akshay Sonawane (apsonawane)]

This pull request improves the ExpandBuffer function in the generation device helper to better support tensors with rank greater than four, and adds a new test to verify this behavior. The main changes are as follows:

Support for high-rank tensors in ExpandBuffer:

• Refactored the ExpandBuffer implementation in generation_device_helper.cc to use a dynamic TensorShapeVector instead of a fixed-size array, allowing support for input tensors with more than four dimensions.

Testing enhancements:

• Added an explicit test, ExpandBufferSupportsRankGreaterThanFour, to beam_search_test.cc to ensure that ExpandBuffer correctly handles tensors with rank greater than four.
• Included the necessary header import for generation_device_helper.h in the test file.

[Copilot]

Pull request overview

This PR updates the CPU generation helper’s ExpandBuffer implementation to handle tensors with rank > 4 by switching from a fixed-size dimension array to a dynamically sized TensorShapeVector, and adds a unit test to validate the new behavior in Beam Search test coverage.

Changes:

• Refactor GenerationCpuDeviceHelper::ExpandBuffer to build the expanded shape from a dynamically sized TensorShapeVector.
• Add a new test ExpandBufferSupportsRankGreaterThanFour to verify shape expansion for 5D inputs.
• Include generation_device_helper.h in the Beam Search test to access ExpandBuffer.

Reviewed changes

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

File	Description
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc	Replaces fixed dims[4] shape storage with a dynamic TensorShapeVector to support high-rank tensors in ExpandBuffer.
onnxruntime/test/contrib_ops/beam_search_test.cc	Adds a focused unit test ensuring ExpandBuffer correctly expands the batch dimension for rank > 4 tensors.

---

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