[Docs] Extend TopK spec with NaN handling info

[Lagmator22]

Details:

Extends the TopK operation specs (TopK-1, TopK-3, TopK-11) with documentation on NaN handling behavior.

What this PR does:

• Documents that NaN ordering in TopK results is implementation-defined, consistent with IEEE 754 semantics
• Notes that different frameworks (NumPy, PyTorch) handle NaN ordering differently
• Provides guidance for users who need deterministic behavior: sanitize NaN values before TopK (e.g. replace with -inf or +inf)

This is a docs-only change with no code modifications. The code fix for NaN handling was reverted per reviewer feedback -- a new TopK version with configurable NaN handling (nan_mode attribute) is planned as a follow-up.

Tickets:

• Related to [Bug]: ArgMax returns incorrect index for all-NaN input #33626

[maxnick]

@nshchego , could you please review?

[Copilot]

Pull request overview

This PR fixes a bug in the CPU reference implementation of TopK where NaN values in the input data would incorrectly appear in the top-K results, blocking valid numbers from being selected.

Changes:

• Added explicit NaN checks in the comparison logic of topk_ref_process to ensure NaN values are treated as smaller than any valid number
• Reformatted the dataFomats vector initialization for improved readability

---

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

[Copilot]

Corrected spelling of 'dataFomats' to 'dataFormats'.

[Copilot]

The reformatting of the dataFomats vector splits related pairs across lines inconsistently. Consider placing each layout pair on its own line for better readability, or keeping all pairs on the same line if they fit within line length limits.

Suggested change

	{LayoutType::nspc, LayoutType::nspc}, {LayoutType::nCsp16c, LayoutType::nCsp16c}, {
	LayoutType::nCsp8c, LayoutType::nCsp8c
	}
	{LayoutType::nspc, LayoutType::nspc},
	{LayoutType::nCsp16c, LayoutType::nCsp16c},
	{LayoutType::nCsp8c, LayoutType::nCsp8c}

[Lagmator22]

Hi @nshchego @maxnick , Ive rebased and cleaned up the PR by removin an accidental formatting change that was unrelated to the fix so now it only contains the nan handling fix (2 lines changed), ready for review.

[Lagmator22]

Hi @maxnick, any update on this? @nshchego hasn't responded yet. would be happy to add tests if needed.

[nshchego]

I believe, need to describe this case in documentation: https://github.com/openvinotoolkit/openvino/tree/master/docs/articles_en/documentation/openvino-ir-format/operation-sets/operation-specs/sort
Compare behavior with TF, PyTorch at least. If they handle it in different ways, add appropriate attribute to operation.
Add some NaNs to test inputs:

• https://github.com/openvinotoolkit/openvino/blob/master/src/tests/functional/base_func_tests/src/base/utils/generate_inputs.cpp#L767
• https://github.com/openvinotoolkit/openvino/blob/master/src/tests/functional/base_func_tests/src/base/utils/generate_inputs.cpp#L782
• https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/topk.cpp#L130

[nshchego]

Is there a plan to fix jit_uni_topk_kernel_f32 as well?

[Lagmator22]

Thanks @nshchego for the detailed review, I've tested and updated the documentation to explicitly clarify that openvino treats NaNs as smaller than valid numbers(consistent with NumPy behavior), which differs from PyTorch (where they are larger).

Regarding the other points:

• Test Cases: I tried injecting NaNs into the existing test inputs, but this causes failures because the test framework compares CPU output against the reference implementation. Since Nan != NaN and the relative order of NaNs is undefined, direct comparison fails even when the values are technically correct. For now, I've relied on the existing tests(which pass) since the reference implementation fix itself is verified.

• JIT Kernel: as expected I analyzed the x64 jit kernel logic and confirmed it uses inconsistent comparison flags (_cmp_lt_os vs _cmp_nle_us) which would handle NaNs differently. Since I'm developing on ARM (Mac) where the JIT path isn't used, I can't verify a jit fix locally. I'd propose addressing the JIT fix in a follow-up PR where we can leverage CI for x64 validation.

Ready for another look and thank u again!

[nshchego]

So, this case is undefined in product at all and we can't simply modify old operations. That may affect other customers.
In my opinion, we should introduce a new operation version. Define this case there and add an attribute to differentiate behavior if there are different approaches in the supported frameworks.

@mitruska, could you please take a look? Do we need a new TopK version for that?

[Lagmator22]

Rebased onto latest master and extended the NaN handling documentation to TopK-1 and TopK-3 specs for consistency with TopK-11 (all 3 versions now describe the same NaN behavior).

@mitruska - I understand this is pending your decision on whether NaN handling should be added to the existing TopK versions or introduced through a new operation version. Happy to go either route. If a new version is the way forward, I will try to help with that as well as a follow up.

In the meantime, the code fix itself is minimal (2 lines in the reference implementation sort path) and all 120 existing smoke_TopK tests pass locally. The JIT kernel fix would be a separate follow-up as discussed with @nshchego.

[mitruska]

Hi @Lagmator22, thank you for contributing to OpenVINO.

This PR can't be merged as a "fix" without common agreement (including all plugins) on the definition of "correct" behavior. IEEE‑754 defines NaN as something that fails ordered comparison. For any real number x (including ±inf) and any NaN, x < NaN, x > NaN, and x == NaN are all false, and even NaN == NaN is false. There’s no standardized and defined position for NaNs in a sort/TopK result, so implementations are free to keep NaNs in place, move NaNs to the end, or otherwise reorder NaNs. This explains the cross-framework/platform differences.

Given that frontend frameworks don’t explicitly document NaN ordering, my recommendation at this point is to keep OpenVINO TopK/ArgMax kernels using current comparisons (for perf) and document NaN ordering as implementation-defined or unspecified rather than adding NaN-specific kernel logic.

If deterministic behavior is needed for a model, it can be sanitized/masked before TopK, e.g. replace NaNs with -inf (push NaNs to bottom) or +inf (push NaNs to top).

Agree with @nshchego that such a change is a candidate for a new version of TopK, having a mode aligning NaN sort options with each frontend. Possible perf impact of additional NaN detection/handling should be measured for different platforms to avoid regressions, as all models using TopK can be affected (even if the NaN is not in the input, it's not possible to determine it at the conversion time for non-const inputs). Perf regression could be a blocker to use new version of TopK even if introduced.

If you are interested in follow-up work, a POC of a new TopK can be prepared to define constraints and benefits of having NaN behavior defined, while it needs to be supported with perf data and clear definition of expected behavior for different formats of models.

[mitruska]

This code is cpu specific, the common reference implementation used for constant folding, testing plugins and results comparison between them is here:

openvino/src/core/reference/include/openvino/reference/topk.hpp

Lines 110 to 125 in 619647f

	* @brief Reference implementation for TopK operator
	*
	* @param arg Pointer to input data.
	* @param out_indices Pointer to output indicies.
	* @param out_values Pointer to output values.
	* @param in_shape Input data shape.
	* @param out_shape Output data (values, indicies) shape.
	* @param axis Axis for search of top K elements.
	* @param k Number to find of top elements.
	* @param compute_max Select mode of find max or min.
	* @param sort Sorting type.
	*/
	template <typename T,
	typename U,
	typename std::enable_if<!std::is_same<typename std::decay<U>::type, int64_t>::value>::type* = nullptr>
	void topk(const T* arg,

[Lagmator22]

@mitruska Thanks for the detailed review, makes complete sense. I'll revert the code change and update the docs to note NaN ordering as implementation-defined behavior instead. If that works for you we can merge this PR as a docs-only change, and then discuss the new TopK version separately.

I'd really like to take on the new TopK version work if you're open to guiding me through it. I have a few questions so I can get started on the right track:

For the new op version, should the NaN handling mode be an attribute (something like nan_mode with options for numpy vs torch style), or is there a different way you'd want it structured?

Where should the reference implementation live, should I be modifying src/core/reference/include/openvino/reference/topk.hpp that you pointed out, or would the new version get its own file?

For perf measurements, what benchmarks would you want to see? I'm thinking comparing the NaN-aware path vs current on different input sizes and NaN densities across CPU at minimum. Any specific setup or tooling you'd recommend?

Is there an existing TopK version bump (like the jump from TopK-3 to TopK-11) I should study as a reference for how new versions are added end to end?

Should I open a separate issue to track the new version?

Would appreciate any pointers you have, really want to get this right.

Thank you

[Lagmator22]

Hi @mitruska,

I've reverted the code change and updated this PR to be docs-only. The three TopK specs (TopK-1, TopK-3, TopK-11) now document NaN ordering as implementation-defined, consistent with your recommendation and IEEE 754 semantics. I also included the -inf/+inf sanitization guidance you suggested for users who need deterministic behavior.

The PR should be straightforward to merge now since it's purely documentation with no code or perf impact.

I'm currently in mid-semester exams (through March 20), but I'm very interested in the follow-up work on a new TopK version with configurable NaN handling modes, happy to start scoping the POC once exams wrap up if you're open to guiding me through it. Also, I am participating in gsoc for the #2 project idea(I have a demo repo as well).

Would appreciate it if you could take a look when you get a chance. Thanks for the guidance!

[mitruska]

Hi @Lagmator22, extending TopK spec with the proposed note reflecting current NaN behavior LGTM.

As this PR no longer provides code changes, please update the PR description.
Any further work should be provided as a separate PR.

---

Regarding new version of TopK POC proposal:

The new op version should be added to the new opset (opset17 need to be initialized), as a new class v17::TopK (reusing TopK Base) and can expose an attribute (enum) to control NaN handling. Suggested attribute: a small enum such as nan_mode with values like nan_as_smallest / nan_as_largest / none (for backward compatibility). For the reference the existing function can be reused / modified (src/core/reference/include/openvino/reference/topk.hpp) as long as the changes are backward compatibile. By default it should preserve current OpenVINO behavior.

Here are example PRs with new operators work:

• #22796
• #28698

With new op proposal please provide motivation of the changes, perfectly with model examples that would benefit from extended TopK NaN handling.

Further review and decisions can be make based on that.

[Lagmator22]

Hi @mitruska, updated the PR description to reflect this is docs-only. @nshchego @kblaszczak-intel could you take a look when you get a chance? Just NaN handling notes added to the TopK-1/3/11 specs, no code changes.

[praasz]

build_jenkins