Is there a Tensorflow equivalent for np.random.choice to randomly sample from a discrete set?

I am trying to implement an image augmentation strategy similar to RandAugment in TensorFlow. From the RandAugment paper, the following code shows how N augmentations are randomly selected to be applied to images.

transforms = [’Identity’, ’AutoContrast’, ’Equalize’, ’Rotate’, ’Solarize’,
              ’Color’,’Posterize’, ’Contrast’, ’Brightness’, ’Sharpness’,
              ’ShearX’, ’ShearY’,’TranslateX’, ’TranslateY’]

def randaugment(N, M):
"""Generate a set of distortions.
Args:
N: Number of augmentation transformations to apply
sequentially.
M: Magnitude for all the transformations.
"""
sampled_ops = np.random.choice(transforms, N)
return [(op, M) for op in sampled_ops]

However, I wish to do this per batch of images in TensorFlow, ideally as efficiently as possible. It would look something like

transform_names = ['Identity', 'Brightness', 'Colour', 'Contrast', 'Equalise', 'Rotate',
                   'Sharpness', 'ShearX', 'ShearY', 'TranslateX', 'TranslateY']

transforms = {'Identity':identity, 'Brightness':brightness, 'Colour':colour,
              'Contrast':contrast, 'Equalise':equalise, 'Rotate':rotate,
              'Sharpness':sharpness, 'ShearX':shear_x, 'ShearY':shear_y,
              'TranslateX':translate_x, 'TranslateY':translate_y}

def brightness(image, M):
    M = tf.math.minimum(M, 0.95)
    M = tf.math.maximum(M, 0.05)
    B = M - 1
    image = tf.image.adjust_brightness(image, delta=B)
    image = tf.clip_by_value(image, clip_value_min=0, clip_value_max=1)

    return image

def augment(image):
    N = 3
    M = tf.random.uniform(minval=0, maxval=1, shape=[])
    sampled_ops = np.random.choice(transform_names, N)

    for op in sampled_ops:
        image = transforms[op](image, M)

    return image

x = tf.data.Dataset.from_tensor_slices(x)
x = x.batch(batch_size)
x_a = x.map(augment)

where x is the dataset of images, and augment is the augmentation function that randomly samples N augmentations to apply to each image. I've added the brightness function to illustrate the composition of the individual augmentation functions. From what I've gathered, any NumPy function seems to only be called once across the entire dataset, meaning the sampled augmentations will be the same for every image.

How could I write this code such that the individual augmentations are randomly sampled independently for each batch?