False positive AllDiff results for models with randomly-initialized Keras layers

[laolvfan]

First, thank you for your excellent work and for sharing it.

Description

Test cases containing randomly-initialized Keras layers (e.g. tf.keras.layers.Dense, Conv2D) are frequently reported as AllDiff, with large differences between the naive and xla runs.

Example

class Model(tf.keras.Model):
    def __init__(self):
        super(Model, self).__init__()
        self.l = tf.keras.layers.Dense(5)
    def call(self, x1):
        v1_0 = tf.nn.relu(x1)
        t = tf.broadcast_to(v1_0, [1, 4])
        x2 = self.l(t)
        return tf.broadcast_to(tf.nn.relu(x2), [1, 5])

m = Model()
input_shape = [4]
x1 = tf.constant([0.0, 1.0, 2.0, 3.0], shape=input_shape)
input_data = [x1]

Reported result

Value mismatch:
[[0.19949746 0. 0. 0.01175135 0. ]] # naive
[[0.25554395 0. 0. 0.29982996 0.83192015]] # xla

What I found

When I manually fix the model weights before running — i.e. build the model once, save the weight values, and restore the same values before each of the naive / xla / autocluster runs — the outputs become identical (differences drop to ~1e-7).

This suggests the three runs are not using the same weights. My guess is that the randomly-initialized Keras layers create their weights at different points (eager vs. tf.function tracing), so even with the same seed the runs end up with different weights, which produces the false AllDiff.