Merge pull request #13525 from tensorflow:LakshmiKalaKadali-patch-4

PiperOrigin-RevId: 733182458

Author: tensorflower-gardener

official/vision/ops/augment.py

@@ -2697,8 +2697,8 @@ def distort(self, images: tf.Tensor,
 
   @staticmethod
   def _sample_from_beta(alpha, beta, shape):
-    sample_alpha = tf.random.gamma(shape, 1., beta=alpha)
-    sample_beta = tf.random.gamma(shape, 1., beta=beta)
+    sample_alpha = tf.random.gamma(shape, alpha, beta=1.0)
+    sample_beta = tf.random.gamma(shape, beta, beta=1.0)
     return sample_alpha / (sample_alpha + sample_beta)
 
   def _cutmix(self, images: tf.Tensor,