tf.keras.activations.sparsemax computes incorrect output for a simple three-element input

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf 2.21.0, also reproduced on tf 2.22.0-dev20260904

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

No response

Python version

Python 3.13.5

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

tf.keras.activations.sparsemax returns an incorrect forward value for a finite float64 input.

In the reproducer below, the input is [[t / 4, 0, -2]] with t = 0.7, so the logits are [[0.175, 0.0, -2.0]]. Near this point, the sparsemax support should contain the first two entries. The threshold is (t / 4 - 1) / 2, so the first output should be t / 4 - threshold = 0.5875.

However, TensorFlow returns 0.5 for the first output component. The result appears to use an incorrect support threshold.

Expected behavior?

The first output component of tf.keras.activations.sparsemax([[0.175, 0.0, -2.0]]) should be close to 0.5875, not 0.5.

Standalone code to reproduce the issue

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["TF_NUM_INTRAOP_THREADS"] = "1"
os.environ["TF_NUM_INTEROP_THREADS"] = "1"

import tensorflow as tf

def target(t):
    x = tf.reshape(
        tf.stack([
            t / tf.constant(4, dtype=t.dtype),
            tf.constant(0, dtype=t.dtype),
            tf.constant(-2, dtype=t.dtype),
        ]),
        (1, 3),
    )
    return tf.keras.activations.sparsemax(x, axis=-1)[0, 0]

x = tf.constant(0.7, dtype=tf.float64)

actual = target(x)
expected = 0.5875

print("actual:", actual.numpy())
print("expected:", expected)

if abs(float(actual) - expected) > 1e-9:
    print("BUG REPRODUCED: tf.keras.activations.sparsemax computes an incorrect output")
else:
    print("not reproduced")

Relevant log output

actual: 0.5
expected: 0.5875
BUG REPRODUCED: tf.keras.activations.sparsemax computes an incorrect output