        def mixup_data_rng(key_0, key_1, x: jnp.ndarray, y: jnp.ndarray, alpha=1):
            """
            performs mixup regularization on the input images and labels
            """
            if alpha > 0:
                lam = jax.random.beta(key_0, alpha, alpha)
            else:
                lam = 1
            batch_size = x.shape[0]
            index = jax.random.permutation(key_1, batch_size)
            mixed_x = lam * x + (1 - lam) * x[index, :]
            y_a, y_b = y, y[index]
            return mixed_x, y_a, y_b, lam
performs mixup regularization on the input images and labels
serl_launcher/serl_launcher/agents/continuous/vice.py