Fix GatherND division by zero when batch dimensions mismatch

[chaya2350]

Fixes #23828

Added validation to check:

• num_batches is not zero
• num_slices is divisible by num_batches

Before this fix, mismatched batch dimensions caused a crash due to division by zero.

Description

This PR fixes a division by zero crash in the GatherND operator when batch dimensions mismatch between input and indices tensors.

Changes made:

Added validation in gather_nd.cc to check that num_batches is not zero before division
Added validation that num_slices is divisible by num_batches
Added a unit test to verify the fix

Motivation and Context

Description
Fixes #23828

When batch_dims is set but the actual batch dimensions of the input tensor and indices tensor don't align correctly, the code performs a division that can result in division by zero, causing a crash.

For example, with:

Input shape: [2, 2, 2]
Indices shape: [2, 1]
batch_dims=1
The calculation num_slices / num_batches would crash if num_batches is 0, or produce unexpected results if they don't divide evenly.

This fix returns a clear error message instead of crashing.

[chaya2350]

@fdwr @justinchuby Could you please review this PR? It fixes the issue you reported. Thanks! 🙂

[justinchuby]

@tianleiwu

[Copilot]

Pull request overview

This PR fixes a division by zero crash in the GatherND operator when batch dimensions between input and indices tensors are mismatched. The fix adds validation to ensure num_batches is non-zero and that num_slices is evenly divisible by num_batches before performing the division.

Changes:

• Added validation checks in gather_nd.cc to prevent division by zero
• Added a unit test to verify the fix handles batch dimension mismatches correctly

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/providers/cpu/tensor/gather_nd.cc	Added validation to check num_batches is not zero and num_slices is divisible by num_batches before division
onnxruntime/test/providers/cpu/tensor/gather_nd_op_test.cc	Added test case to verify batch dimension mismatch returns an error instead of crashing

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[chaya2350]

I've addressed the Copilot comment. Ready for re-review! 🙂

[Copilot]

Pull request overview

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

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

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

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[fdwr]

👍 Thanks :).

[justinchuby]

Why was this closed?

[fdwr]

@justinchuby - Weird, the branch was deleted too.

[chaya2350]

This was a mistake...
I got this one confused with another PR
Apologies!

[justinchuby]

@chaya2350 could you run lintrunner f to format the code? Thanks

[Copilot]

Pull request overview

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

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tianleiwu]

@chaya2350, please format the code with lintrunner -a. The "Lint / Python format" check is required to pass for pull request.
See https://github.com/microsoft/onnxruntime/blob/main/docs/Coding_Conventions_and_Standards.md#linting

LGTM. The changes are safe and improve the robustness of the GatherND operator on CPU.

[justinchuby]

Looks like there are some errors with CUDA. @tianleiwu do you have suggestions?

[chaya2350]

Hi,
The failing test is intended to run only on the CPU execution provider.
As defined in the test code, I explicitly excluded all other EPs (including CUDA) using the exclude_providers argument in OpTester::Run.
However, according to the test log, the test is still being executed with CUDAExecutionProvider:

registered execution providers: CUDAExecutionProvider
This suggests that the exclusion is not being respected in the CI environment, or that the test runner is forcing all tests to run on all available EPs.
Could you please advise how to ensure this test is executed only on CPU as intended?
If there is a CI configuration or test runner setting that needs to be updated, I would appreciate your guidance.
Thank you!

[justinchuby]

Based on the team's suggestions, could you take a look at

onnxruntime/onnxruntime/core/providers/cuda/tensor/gather_nd.cc

Line 30 in 1a71a5f

ORT_RETURN_IF_NOT(

and ensure the error strings match with the cpu provider?

[tianleiwu]

@chaya2350, test can explicitly run on CPU provider only like

    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
    execution_providers.push_back(DefaultCpuExecutionProvider());
    test.Run(OpTester::ExpectResult::kExpectFailure,
           "GatherND: input tensor batch dimensions cannot be zero", {}, nullptr, &execution_providers);

see

onnxruntime/onnxruntime/test/unittest_util/base_tester.h

Lines 576 to 581 in 735e69a

	void Run(ExpectResult expect_result = ExpectResult::kExpectSuccess, const std::string& expected_failure_string = "",
	const std::unordered_set<std::string>& excluded_provider_types = {},
	const RunOptions* run_options = nullptr,
	std::vector<std::unique_ptr<IExecutionProvider>>* execution_providers = nullptr,
	ExecutionMode execution_mode = ExecutionMode::ORT_SEQUENTIAL,
	const Graph::ResolveOptions& resolve_options = {});