Honour keepdims when axis is None in the torch backend reductions by Sid-5137 · Pull Request #23535 · keras-team/keras

Sid-5137 · 2026-08-30T18:14:50Z

Description

On the torch backend, keepdims=True is silently ignored when axis=None for sum, prod, max, min, amax and amin. A (2, 3) input returns a scalar instead of shape (1, 1), so code that relies on the rank being preserved gets the wrong shape with no error.

import numpy as np
from keras import ops

x = np.arange(6, dtype="float32").reshape(2, 3)
for name in ("sum", "prod", "max", "min", "amax", "amin",
             "mean", "std", "var", "any", "all"):
    out = ops.convert_to_numpy(
        getattr(ops, name)(ops.convert_to_tensor(x), axis=None, keepdims=True)
    )
    print(name, out.shape)

                                   torch      numpy / jax
sum, prod, max, min, amax, amin    ()         (1, 1)
mean, std, var, any, all           (1, 1)     (1, 1)

numpy and jax are already correct, and so are mean, std, var, any and all on the torch backend itself, which is what makes this an inconsistency rather than an undocumented limitation.

Cause

Each of the six took a separate axis is None branch that never forwarded keepdim:

# sum
if axis is not None:
    return cast(torch.sum(x, axis=axis, keepdim=keepdims), dtype)
return cast(torch.sum(x), dtype)          # keepdims dropped

# amax / amin
if axis is None:
    return torch.amax(x)                  # keepdims dropped

torch.sum, torch.amax and torch.amin all accept dim=None and honour keepdim correctly, so those branches are simply removed.

torch.prod is the exception, its dim must be an int and it rejects dim=None, so it reduces everything and then restores the rank when keepdims=True.

max and min had their own bypass calling torch.max / torch.min directly instead of going through amax / amin, so fixing amax / amin alone did not fix them. They now route through those.

Tests

Adds NumpyOneInputOpsCorrectnessTest::test_reductions_keepdims_with_axis_none, covering all six ops against numpy. The existing tests only ever exercise an explicit axis, which is why this was not caught.

Verified the test actually catches the bug: reverting just the backend change and keeping the test gives

AssertionError: Tuples differ: () != (1, 1)
+ (1, 1) : sum(axis=None, keepdims=True) shape mismatch

On the openvino backend the four affected assertions are guarded inline with a TODO(#23536): that backend has its own, separate axis=None + keepdims defects and cannot pass them yet. This PR only touches the torch backend, so fixing those is out of scope here, and the guards can be dropped once #23536 is resolved.

Verification

Full keras/src/ops suite:

backend	result
torch	5525 passed, 752 skipped
numpy	6978 passed, 752 skipped
jax	7590 passed, 140 skipped
tensorflow	6954 passed, 38 skipped

ruff check and ruff format --check are clean.

AI assistance disclosure

Per the AI-Assisted Contribution Policy: I used Claude Code on this change. It found the inconsistency by comparing reduction behaviour across backends, traced each of the separate axis is None branches, drafted the fix and the regression test, and ran the suites locally. I reviewed the change, I understand why torch.prod needs different handling from the others and why max / min needed fixing separately, and I will be answering review comments myself.

Contributor Agreement

I am a human, and not a bot.
I will be responsible for responding to review comments in a timely manner.
I will work with the maintainers to push this PR forward until submission.

gemini-code-assist

Code Review

This pull request simplifies several reduction operations (max, min, amax, amin, sum) in the PyTorch backend by removing redundant axis checks, and fixes prod to correctly handle keepdims=True when axis=None. A new test suite is also added to verify rank preservation for these reductions. The feedback suggests converting axis=None to a tuple of all dimensions in amax and amin to prevent potential PyTorch overload resolution issues when dim=None is passed.

gemini-code-assist · 2026-08-30T18:16:28Z

    x = convert_to_tensor(x)
-    if axis is None:
-        return torch.amax(x)
    if axis == () or axis == []:


In PyTorch, passing dim=None to torch.amax can be problematic because of PyTorch's overload resolution rules (where dim=None may not be accepted or may cause keepdim to be ignored in some PyTorch versions). Converting axis=None to a tuple of all dimensions (tuple(range(x.ndim))) is a more robust approach that is universally supported across all PyTorch versions and guarantees that the correct multi-dimensional reduction overload is matched.

Suggested change

x = convert_to_tensor(x)

if axis is None:

return torch.amax(x)

if axis == () or axis == []:

x = convert_to_tensor(x)

if axis is None:

axis = tuple(range(x.ndim))

if axis == () or axis == []:

gemini-code-assist · 2026-08-30T18:16:29Z

    x = convert_to_tensor(x)
-    if axis is None:
-        return torch.amin(x)
    if axis == () or axis == []:


In PyTorch, passing dim=None to torch.amin can be problematic because of PyTorch's overload resolution rules (where dim=None may not be accepted or may cause keepdim to be ignored in some PyTorch versions). Converting axis=None to a tuple of all dimensions (tuple(range(x.ndim))) is a more robust approach that is universally supported across all PyTorch versions and guarantees that the correct multi-dimensional reduction overload is matched.

Suggested change

x = convert_to_tensor(x)

if axis is None:

return torch.amin(x)

if axis == () or axis == []:

x = convert_to_tensor(x)

if axis is None:

axis = tuple(range(x.ndim))

if axis == () or axis == []:

On the torch backend, `keepdims=True` was silently ignored when `axis=None` for sum, prod, max, min, amax and amin. A (2, 3) input returned a scalar instead of shape (1, 1), so downstream code that relies on the rank being preserved got the wrong shape with no error. The numpy and jax backends already behave correctly, and so do mean, std, var, any and all on the torch backend itself, which is what made this inconsistent rather than merely undocumented. torch, axis=None, keepdims=True, input (2, 3): sum, prod, max, min, amax, amin -> () before mean, std, var, any, all -> (1, 1) Each function took a separate `axis is None` branch that dropped keepdim. torch.sum, torch.amax and torch.amin all accept dim=None and handle keepdim correctly, so those branches are simply removed. torch.prod does not accept dim=None, so it reduces everything and then restores the rank. max and min had their own bypass calling torch.max / torch.min directly rather than going through amax / amin, so they now route through those. Adds a regression test covering all six ops. It fails before this change with `Tuples differ: () != (1, 1)`.

codecov-commenter · 2026-08-30T18:35:53Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.29%. Comparing base (78d4401) to head (f0e3a1f).
⚠️ Report is 11 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #23535      +/-   ##
==========================================
- Coverage   85.02%   84.29%   -0.74%     
==========================================
  Files         468      468              
  Lines       71087    71141      +54     
  Branches    11788    11788              
==========================================
- Hits        60442    59966     -476     
- Misses       7632     8174     +542     
+ Partials     3013     3001      -12

Flag	Coverage Δ
keras	`84.11% <100.00%> (-0.72%)`	⬇️
keras-cpu	`84.11% <100.00%> (+0.01%)`	⬆️
keras-gpu	`?`
keras-jax	`58.33% <0.00%> (-0.30%)`	⬇️
keras-numpy	`53.98% <0.00%> (+0.03%)`	⬆️
keras-openvino	`59.68% <0.00%> (+0.02%)`	⬆️
keras-tensorflow	`59.97% <0.00%> (-0.26%)`	⬇️
keras-torch	`59.37% <100.00%> (-0.37%)`	⬇️
keras-tpu	`?`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Addresses review feedback: older torch versions do not accept dim=None for torch.amax / torch.amin, so reduce over an explicit tuple of all dimensions instead. Behaviour on current torch is unchanged, verified against numpy for 0-d, 1-d and 2-d inputs.

The openvino backend has its own axis=None + keepdims=True defects (sum, prod, amax, amin, std, var, any and all return shape (1,) instead of preserving the rank, and mean raises AttributeError), so the new test cannot pass there yet. Excluded via the existing openvino skip list; the backend gap is reported separately.

hertschuh

Thanks for catching this!

hertschuh · 2026-09-01T21:19:32Z

    if isinstance(getattr(result, "values", None), torch.Tensor):
        result = result.values


This is no longer needed, it was an artifact of using torch.max I believe.

hertschuh · 2026-09-01T21:20:28Z

    if isinstance(getattr(result, "values", None), torch.Tensor):
        result = result.values


hertschuh · 2026-09-01T21:24:04Z

+    def test_reductions_keepdims_with_axis_none(self):
+        # `keepdims=True` must preserve the rank even when `axis=None`,
+        # matching numpy. The tests above only cover an explicit axis.
+        x = np.array([[1, 2, 3], [3, 2, 1]])
+        for op_name in ("sum", "prod", "max", "min", "amax", "amin"):
+            knp_op = getattr(knp, op_name)
+            np_op = getattr(np, op_name)
+            expected = np_op(x, axis=None, keepdims=True)
+            result = knp_op(x, axis=None, keepdims=True)
+            self.assertEqual(
+                tuple(result.shape),
+                expected.shape,
+                msg=f"{op_name}(axis=None, keepdims=True) shape mismatch",
+            )
+            self.assertAllClose(result, expected)


Instead of this, please add one test case in each one of the existing test functions for each op. Don't add a new test function just add a line in the existing one.

The reason is that we're going to split tests by op and tests like this one won't play nice with this refactor.

Thanks!

…functions - max and min: the .values unwrapping was only ever needed for torch.max(x, dim=...), which returns a namedtuple. amax and amin always return a plain tensor, so the branch was unreachable. Removed, and both now route straight through amax / amin. - The combined test_reductions_keepdims_with_axis_none is gone. Each of the existing per-op tests (sum, prod, max, min, amax, amin) now covers axis=None with keepdims=True inline. assertAllClose broadcasts a scalar against (1, 1) and would pass on the broken code, so each op asserts the output shape explicitly first. Verified the tests fail without the backend fix and pass with it. - The openvino skip moved from excluded_concrete_tests.txt to inline backend guards on the four ops openvino gets wrong, with a TODO pointing at keras-team#23536. max and min pass on openvino and are unguarded. excluded_concrete_tests.txt is back to untouched.

Sid-5137 · 2026-09-02T04:25:15Z

Thanks for the review! All three changes are done.

Removed the .values unwrapping for both max and min. I checked why it was there before deleting it: only torch.max(x, dim=...) returns a namedtuple with .values, whereas torch.amax and torch.amin always return a plain tensor. Since both functions now go through amax/amin, that branch could never be hit. You were right that it was leftover from torch.max.
Removed the combined test function. Each of the six existing per-op tests now covers axis=None and keepdims=True inline. While moving them, I noticed that a plain assertAllClose passes even with the broken code because it broadcasts the scalar against (1, 1). So each op now explicitly asserts the output shape first, followed by the values. I verified that the tests fail without the backend fix and pass with it.
The four assertions that OpenVINO can't pass yet are now guarded inline with a TODO referencing keepdims=True with axis=None does not preserve rank on the openvino backend #23536. excluded_concrete_tests.txt is back to untouched.

google-ml-butler Bot added the size:S label Aug 30, 2026

google-ml-butler Bot assigned gbaned Aug 30, 2026

gemini-code-assist Bot reviewed Aug 30, 2026

View reviewed changes

Sid-5137 force-pushed the fix/torch-keepdims-axis-none branch from ab8d055 to e6fc36b Compare August 30, 2026 18:23

Sid-5137 added 2 commits August 31, 2026 00:05

Sid-5137 mentioned this pull request Aug 30, 2026

keepdims=True with axis=None does not preserve rank on the openvino backend #23536

Open

hertschuh requested changes Sep 1, 2026

View reviewed changes

hertschuh added the stat:awaiting response from contributor label Sep 1, 2026

google-ml-butler Bot removed the stat:awaiting response from contributor label Sep 2, 2026

Sid-5137 requested a review from hertschuh September 2, 2026 04:25

google-ml-butler Bot added the awaiting review label Sep 2, 2026

Navigation Menu

Uh oh!

Honour keepdims when axis is None in the torch backend reductions - #23535

Honour keepdims when axis is None in the torch backend reductions#23535
Sid-5137 wants to merge 4 commits into
keras-team:masterfrom
Sid-5137:fix/torch-keepdims-axis-none

Sid-5137 commented Aug 30, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 30, 2026

Uh oh!

gemini-code-assist Bot Aug 30, 2026

Uh oh!

codecov-commenter commented Aug 30, 2026 •

edited

Loading

Uh oh!

hertschuh left a comment

Uh oh!

hertschuh Sep 1, 2026

Uh oh!

hertschuh Sep 1, 2026

Uh oh!

hertschuh Sep 1, 2026

Uh oh!

Sid-5137 commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		if isinstance(getattr(result, "values", None), torch.Tensor):
		result = result.values

Uh oh!

Conversation

Sid-5137 commented Aug 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Cause

Tests

Verification

AI assistance disclosure

Contributor Agreement

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 30, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 30, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Aug 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

hertschuh Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

hertschuh Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

hertschuh Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

Sid-5137 commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Sid-5137 commented Aug 30, 2026 •

edited

Loading

codecov-commenter commented Aug 30, 2026 •

edited

Loading