MixUp + Cutmix implementation is (badly!) incorrect

[dibyaghosh]

The code for sampling from a beta distribution in tfm.vision.augment.MixupAndCutmix._sample_from_beta is incorrect:

It should be

  @staticmethod
  def _sample_from_beta(alpha, beta, shape):
    sample_alpha = tf.random.gamma(shape, alpha, beta=1.0) # not: tf.random.gamma(shape, 1.0, beta=alpha)
    sample_beta = tf.random.gamma(shape, beta, beta=1.0) # not: tf.random.gamma(shape, 1.0, beta=beta)
    return sample_alpha / (sample_alpha + sample_beta)

https://github.com/tensorflow/models/blob/8c5c79c27c4bb2843308bf8623daa24f4f3cbe62/official/vision/ops/augment.py#L2712C1-L2713C1

This results in very bad mixup samples:

# code for plot
import tensorflow_models as tfm
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
a = tf.range(1000)
x = tfm.vision.augment.MixupAndCutmix._sample_from_beta(0.2, 0.2, tf.shape( tf.range(1000))).numpy()
sns.kdeplot(x, clip=(0, 1))
x = np.random.beta(0.2, 0.2, 100000)
sns.kdeplot(x, clip=(0, 1))
plt.show()

[LakshmiKalaKadali]

Hi @dibyaghosh,

Thanks for bringing out the issue and raising PR. I have reproduced the issue and observed the deviation between beta distribution in tfm.vision.augment.MixupAndCutmix._sample_from_beta and np.random.beta(0.2, 0.2, 100000) . Here is the [gist].(https://colab.sandbox.google.com/gist/LakshmiKalaKadali/06533824610d6e85ea4aa3c6399819e6/tf_model_13490.ipynb). The observations are attached:

Please track the PR raised for this issue as well.

Thank You