def shake_drop_eval(x, mask_prob, alpha_min, alpha_max):
  """ShakeDrop eval pass.

  See https://arxiv.org/abs/1802.02375

  Args:
    x: input to apply ShakeDrop to
    mask_prob: mask probability
    alpha_min: alpha range lower
    alpha_max: alpha range upper

  Returns:
  """
  expected_alpha = (alpha_max + alpha_min) / 2
  # See Eqn 6 in https://arxiv.org/abs/1802.02375
  x = (mask_prob + expected_alpha - mask_prob * expected_alpha) * x
  return x