Repository: tensorflow
Issue: MLP subsequence is inconsistent with full sequence
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.15

### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

v1.12.1-105526-gf9a553f735c 2.16.0-dev20240206

### Custom code

No

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

**Issue:** Given a two layer MLP with a reshape, I notice that `model(data)[:, :end_idx])` does not equal `model(data[:, :end_idx])`.
That is, comparing the output of the model on a subsequence of the data is _not_ the same as the corresponding subsequence of the output, even when the MLP should not have any interactions across the time-dimension.

Can you confirm this is expected behaviour? I notice that GPUs (T4 on colab) has more inconsistencies than CPU, so I assume this is some sort of computional optimizations. 
Is there anyway to guarantee that the outputs match? I have tried  `tf.config.experimental.enable_op_determinism()`.

Thanks!
Erik

### Standalone code to reproduce the issue

```shell
Colab: https://colab.research.google.com/drive/1ffw4XhJKocu71ORANA_70eG7Tl20-cLh?usp=sharing


Code from colab:

import tensorflow as tf

tf.keras.utils.set_random_seed(1)
tf.config.experimental.enable_op_determinism()

class BrokenModel(tf.keras.Model):

    def __init__(self,):
        super().__init__()

        self.dense1 = tf.keras.layers.Dense(
            32, activation=None,
        )
        self.dense2 = tf.keras.layers.Dense(
            256 , activation=None,
        )

    def call(self, x: tf.Tensor):
        dims = tf.shape(x)

        # First dense
        x = self.dense1(x)

        # Flatten last two dimensions
        x = tf.reshape(
            x, (dims[0], dims[1],  -1)
        )

        # Second dense
        x = self.dense2(x)
        return x

def evaluate_subsequence_consistency(
    model,
    data,
):

    orig_output = model(data)

    for end_idx in range(0, data.shape[1]):
        # Slice subsequence
        output = model(data[:, :end_idx])

        # Check if output is the same
        res_str = ""
        for j in range(0, end_idx):
            # Assumes output has shape [B, T, ...]
            # Check if each time step is identical between subsequence and full sequence
            orig_data = orig_output[:, j].numpy()
            new_data = output[:, j].numpy()

            if tf.experimental.numpy.allclose(orig_data, new_data):
                res_str += "."
            else:
                res_str += "X"

        print(f"model(data)[:, :{end_idx}] vs model(data[:, :{end_idx}]) => {res_str}")

model = BrokenModel()
data = tf.ones((1, 313, 33, 3))
evaluate_subsequence_consistency(model, data)
```


### Relevant log output

```shell
X = output does not match, . = output is matching

model(data)[:, :1] vs model(data[:, :1]) => X
model(data)[:, :2] vs model(data[:, :2]) => XX
model(data)[:, :3] vs model(data[:, :3]) => XXX
model(data)[:, :4] vs model(data[:, :4]) => XXXX
model(data)[:, :5] vs model(data[:, :5]) => XXXXX
model(data)[:, :6] vs model(data[:, :6]) => XXXXXX
model(data)[:, :7] vs model(data[:, :7]) => XXXXXXX
model(data)[:, :8] vs model(data[:, :8]) => XXXXXXXX
model(data)[:, :9] vs model(data[:, :9]) => XXXXXXXXX
model(data)[:, :10] vs model(data[:, :10]) => XXXXXXXXXX
model(data)[:, :11] vs model(data[:, :11]) => XXXXXXXXXXX
model(data)[:, :12] vs model(data[:, :12]) => XXXXXXXXXXXX
model(data)[:, :13] vs model(data[:, :13]) => XXXXXXXXXXXXX
model(data)[:, :14] vs model(data[:, :14]) => XXXXXXXXXXXXXX
model(data)[:, :15] vs model(data[:, :15]) => XXXXXXXXXXXXXXX
model(data)[:, :16] vs model(data[:, :16]) => XXXXXXXXXXXXXXXX
model(data)[:, :17] vs model(data[:, :17]) => XXXXXXXXXXXXXXXXX
model(data)[:, :18] vs model(data[:, :18]) => XXXXXXXXXXXXXXXXXX
model(data)[:, :19] vs model(data[:, :19]) => XXXXXXXXXXXXXXXXXXX
model(data)[:, :20] vs model(data[:, :20]) => XXXXXXXXXXXXXXXXXXXX
```