TF2 Data Augmentation Randomness

I've been trying to add custom augmentation to a model. However, while library functions such as tf.image.random_flip_left_right or tf.image.random_brightness actually work and create random variations for every image in every sample, my custom functions always have the same effect for ALL images in ALL batches.

For instance, these are my functions for "random" gaussian and motion blurs:

@tf.function
def gaussian(image, n_channels, dtype):
  size = np.random.randint(2, 6)
  sigma = np.random.uniform(0.5, 4.0)
  x = tf.range(-size // 2 + 1, size // 2 + 1, dtype=dtype)
  g = tf.math.exp(-(tf.pow(x, 2) / (2 * tf.pow(tf.cast(sigma, dtype), 2))))
  g_norm2d = tf.pow(tf.reduce_sum(g), 2)
  k = tf.tensordot(g, g, axes=0) / g_norm2d
  k = tf.expand_dims(k, axis=-1)
  k = tf.expand_dims(tf.tile(k, (1, 1, n_channels)), axis=-1)
  return tf.nn.depthwise_conv2d(image[None], k, [1,1,1,1], 'SAME')[0] 

@tf.function
def motion(image, n_channels, dtype):
    size = np.random.randint(2, 11)
    k = np.zeros((size, size))
    k[int((size-1)/2), :] = np.ones(size)
    k = k / size
    k = tf.convert_to_tensor(k, dtype=dtype)
    k = tf.expand_dims(k, axis=-1)
    k = tf.expand_dims(tf.tile(k, (1, 1, n_channels)), axis=-1)
    return tf.nn.depthwise_conv2d(image[None], k, [1,1,1,1], 'SAME')[0] 

As I said, my images get random brightness, flip, etc. But they ALL get the same motion and Gaussian blur, which is not what I wanted.

EDIT: To clarify, I call all augmentation functions sequentially in an augment(image, label) function, which is called via dataset.map(augment). This works just fine for the other augmentations.