Select subset of weights from a list of keras layers in @tf.function

Question

I have a @tf.function decorated function and I would like to select a random subset of weights for each layer of a keras NN. The function takes as input a list of weight var