Expected Randomness does not occur in Tensorflow Layer

Question

I wrote a custom layer that shuffle the input. When I try to test the layer, said shuffling does not occur. Here is my minimal noise layer below:

class ShuffleLayer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(ShuffleLayer, self).__init__(**kwargs)
        
    def call(self, inputs, training=None):
        if training:
            shuffled = tf.stop_gradient(tf.random.shuffle(inputs))
            return shuffled
        return inputs

When I test the layer, the layer will not shuffle

SL = ShuffleLayer()
x = tf.reshape(tf.range(0,10, dtype=tf.float32), (5,2))
y = SL(x)
print(x.numpy())
print(y.numpy())

[[0. 1.] [2. 3.] [4. 5.] [6. 7.] [8. 9.]]

[[0. 1.] [2. 3.] [4. 5.] [6. 7.][8. 9.]]

Why will the expected behavior not occur?

Answer 1

Looking at the layer call, we see that the layer does nothing if training is None. When the layer is called as y = SL(x), it sees that training is None and returns the inputs. Getting the shuffled output is done by turning on the training parameter:

y = SL(x, training=True)
print(x.numpy())
print(y.numpy())

[[0. 1.][2. 3.][4. 5.][6. 7.][8. 9.]]

[[0. 1.][6. 7.][2. 3.][8. 9.][4. 5.]]