Repository: tensorflow
Issue: Randomization generating repeated sequences with `tf.function`, `tf.random.set_seed` and `tf.cond`
Labels: stat:awaiting tensorflower,type:bug,comp:ops

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

2.15.0-dev20230919

### Custom code

No

### OS platform and distribution

Linux Ubuntu 22.04.2 LTS

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

We have encountered an unexpected issue with `tf.function` compilation combined with the global/graph-level seed setting (`tf.random.set_seed`); under some circumstances the random variables are generating identical sequences within the same function. The issue seems to be connected with TensorFlow conditionals (`tf.cond`) within the function, particularly when the branches contain the randomization calls. Removing `tf.cond` results in the generation of unique random values, as expected.

We understand the following expected behavior under `tf.function` compilation when a global/graph-level seed is set, as per the [TensorFlow documentation](https://www.tensorflow.org/api_docs/python/tf/random/set_seed):

> Note that `tf.function` acts like a re-run of a program in this case. When the global seed is set but operation seeds are not set, the sequence of random numbers are the same for each `tf.function`.

However, the problem we are observing arises within a single instance of the function, not across multiple instances. The behavior is as if the internal counters
in `tf.random.uniform` get reset when there is a TensorFlow conditional in the graph. Is this expected behavior?

To illustrate this issue, the example below presents two scenarios. The first one ("SAD" mode) shows the function producing repeated random sequences when `tf.cond` is present. The second scenario ("HAPPY" mode) shows the function generating unique random values once `tf.cond` is removed.


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from pprint import pprint

tf.random.set_seed(42)

var = tf.Variable(0, dtype=tf.int32)
def get_value(happy):
    if happy:
        return tf.stack([tf.random.uniform(()) for _ in range(2)])
    else:
        return tf.cond(
            var == 0,
            lambda: tf.stack([tf.random.uniform(()) for _ in range(2)]),
            lambda: tf.stack([tf.random.uniform(()) for _ in range(2)]),
        )

def randomize(happy):
    return [get_value(happy) for _ in range(3)]

print("SAD")
pprint(tf.function(randomize)(False))
print("HAPPY")
pprint(tf.function(randomize)(True))
```


### Relevant log output

```shell
SAD
[<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.8354591 , 0.15012848], dtype=float32)>,
 <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.8354591 , 0.15012848], dtype=float32)>,
 <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.8354591 , 0.15012848], dtype=float32)>]
HAPPY
[<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.8354591 , 0.15012848], dtype=float32)>,
 <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.98781276, 0.63789964], dtype=float32)>,
 <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.00857747, 0.02621067], dtype=float32)>]
```