keras/src/backend/tensorflow/random.py/shuffle
def shuffle(x, axis=0, seed=None):
    from keras.src.backend.tensorflow.numpy import swapaxes

    seed = _cast_seed(draw_seed(seed))
    if axis == 0:
        return tf.random.experimental.stateless_shuffle(x, seed=seed)
    x = swapaxes(x, axis1=0, axis2=axis)
    x = tf.random.experimental.stateless_shuffle(x, seed=seed)
    x = swapaxes(x, axis1=0, axis2=axis)
    return x