def METHOD_NAME(self, images, transformations, **kwargs):
    batch_size = tf.shape(images)[0]
    height, width = images.shape[1], images.shape[2]
    num_channels = images.shape[3]
    indices = transformations
    # append batch indexes next to shuffled indices
    batch_indexs = tf.repeat(tf.range(batch_size), self.groups)
    batch_indexs = tf.reshape(batch_indexs, (batch_size, self.groups))
    indices = tf.stack([batch_indexs, indices], axis=-1)
    if not num_channels % self.groups == 0:
        raise ValueError(
            "The number of input channels should be "
            "divisible by the number of groups."
            f"Received: channels={num_channels}, groups={self.groups}"
        )
    channels_per_group = num_channels // self.groups
    images = tf.reshape(
        images, [batch_size, height, width, self.groups, channels_per_group]
    )
    images = tf.transpose(images, perm=[0, 3, 1, 2, 4])
    images = tf.gather_nd(images, indices=indices)
    images = tf.transpose(images, perm=[0, 2, 3, 4, 1])
    images = tf.reshape(images, [batch_size, height, width, num_channels])
    return images