fix(ops): Handle floating-point edge cases in ops.argmax()

keras/src/backend/jax/numpy.py

@@ -353,7 +353,14 @@ def arctanh(x):
 
 
 def argmax(x, axis=None, keepdims=False):
-    return jnp.argmax(x, axis=axis, keepdims=keepdims)
+    if x.ndim == 0:
+        return jnp.argmax(x, axis=axis, keepdims=keepdims)
+    x_float = x.astype(jnp.float32)
+    is_negative_zero = (x_float == 0.0) & jnp.signbit(x_float)
+    x_adjusted = jnp.where(
+        is_negative_zero, -jnp.finfo(x_float.dtype).tiny, x_float
+    )
+    return jnp.argmax(x_adjusted, axis=axis, keepdims=keepdims)
 
 
 def argmin(x, axis=None, keepdims=False):

keras/src/backend/numpy/numpy.py

@@ -245,8 +245,14 @@ def arctanh(x):
 
 
 def argmax(x, axis=None, keepdims=False):
-    axis = standardize_axis_for_numpy(axis)
-    return np.argmax(x, axis=axis, keepdims=keepdims).astype("int32")
+    if x.ndim == 0:
+        return np.argmax(x, axis=axis, keepdims=keepdims).astype("int32")
+    x_float = x.astype(np.float32)
+    is_negative_zero = (x_float == 0.0) & np.signbit(x_float)
+    x_adjusted = np.where(
+        is_negative_zero, -np.finfo(x_float.dtype).tiny, x_float
+    )
+    return np.argmax(x_adjusted, axis=axis, keepdims=keepdims).astype("int32")
 
 
 def argmin(x, axis=None, keepdims=False):

keras/src/backend/tensorflow/numpy.py

@@ -837,12 +837,39 @@ def _keepdims(x, y, axis):
 
 
 def argmax(x, axis=None, keepdims=False):
-    _x = x
+    x_float = tf.cast(x, tf.float32)
+    is_negative_zero = tf.logical_and(
+        tf.equal(x_float, 0.0),
+        tf.less(
+            tf.bitwise.bitwise_and(
+                tf.bitcast(x_float, tf.int32),
+                # tf.float32 sign bit
+                tf.constant(0x80000000, dtype=tf.int32),
+            ),
+            0,
+        ),
+    )
+    non_zero_mask = tf.not_equal(x_float, 0.0)
+    masked_abs = (
+        tf.abs(x_float)
+        + (1.0 - tf.cast(non_zero_mask, tf.float32)) * tf.float32.max
+    )
+    min_non_zero = tf.reduce_min(masked_abs) - 1e-9
+    x_adjusted = tf.where(is_negative_zero, -min_non_zero, x_float)
     if axis is None:
-        x = tf.reshape(x, [-1])
-    y = tf.argmax(x, axis=axis, output_type="int32")
-    if keepdims:
-        y = _keepdims(_x, y, axis)
+        x_adjusted = tf.reshape(x_adjusted, [-1])
+        y = tf.argmax(x_adjusted, axis=0, output_type=tf.int32)
+        if keepdims:
+            y = tf.reshape(y, [1, 1])
+    else:
+        rank = tf.rank(x_adjusted)
+        axis_tensor = tf.convert_to_tensor(axis, dtype=tf.int32)
+        positive_axis = tf.cond(
+            axis_tensor < 0, lambda: axis_tensor + rank, lambda: axis_tensor
+        )
+        y = tf.argmax(x_adjusted, axis=positive_axis, output_type=tf.int32)
+        if keepdims:
+            y = tf.expand_dims(y, axis=positive_axis)
     return y
 
 

keras/src/ops/numpy_test.py

@@ -1142,6 +1142,16 @@ def test_argmax(self):
         self.assertEqual(knp.argmax(x, axis=1).shape, (None, 3))
         self.assertEqual(knp.argmax(x, keepdims=True).shape, (None, 3, 3))
 
+    @pytest.mark.skipif(
+        keras.config.backend() == "openvino",
+        reason="OpenVINO doesn't support this change",
+    )
+    def test_argmax_negative_zero(self):
+        input_data = np.array(
+            [-1.0, -0.0, 1.401298464324817e-45], dtype=np.float32
+        )
+        self.assertEqual(knp.argmax(input_data), 2)
+
     def test_argmin(self):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.argmin(x).shape, ())