code: def multiplicative_jitter(x, epsilon=1e-2):
    
    if epsilon == 0:
      return x
    return x * mtf.random_uniform(
        x.mesh, x.shape, minval=1.0 - epsilon, maxval=1.0+epsilon, dtype=x.dtype)
docs: Multiply values by a random number between 1-epsilon and 1+epsilon.
  
    Makes models more resilient to rounding errors introduced by bfloat16.
    This seems particularly important for logits.
  
    Args:
      x: a mtf.Tensor
      epsilon: a floating point value
  
    Returns:
      a mtf.Tensor with the same type and shape as x.