Arithmetic optimization changes tf.argmin result after ReLU tie case

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.21.0

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

No response

Python version

3.11

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

TensorFlow returns a different result for tf.argmin(tf.nn.relu(x)) when arithmetic optimization is enabled.

The input contains several negative values. After ReLU, multiple elements become 0, so argmin should return the first minimum index. Eager execution returns [1], and graph execution with arithmetic_optimization disabled also returns [1]. However, with arithmetic_optimization enabled, the result becomes [2].

In the reproducer below, arithmetic optimization changes the argmin result from [1] to [2].

Expected behavior?

Arithmetic optimization should not change the result of tf.argmin(tf.nn.relu(x)). The optimized graph should preserve the same tie-breaking behavior as eager execution.

Standalone code to reproduce the issue

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf


x = tf.constant([[3.0, -2.0, -7.0, 4.0, -1.0]])

eager = tf.argmin(tf.nn.relu(x), axis=-1).numpy().tolist()

tf.config.optimizer.set_experimental_options({"arithmetic_optimization": True})
on = tf.function(lambda v: tf.argmin(tf.nn.relu(v), axis=-1))(x).numpy().tolist()

tf.config.optimizer.set_experimental_options({"arithmetic_optimization": False})
off = tf.function(lambda v: tf.argmin(tf.nn.relu(v), axis=-1))(x).numpy().tolist()

print(f"eager: {eager}")
print(f"arith on : {on}")
print(f"arith off: {off}")

if on != off:
    print("BUG REPRODUCED: arithmetic optimization changes argmin result")
else:
    print("not reproduced")

Relevant log output

eager: [1]
arith on : [2]
arith off: [1]
BUG REPRODUCED: arithmetic optimization changes argmin result