def shake_shake_train()

in jcm/models/wideresnet_noise_conditional.py [0:0]

• 14 lines of code
• 6 McCabe index (conditional complexity)

def shake_shake_train(xa, xb, rng=None):
    """Shake-shake regularization in training mode.

    Shake-shake regularization interpolates between inputs A and B
    with *different* random uniform (per-sample) interpolation factors
    for the forward and backward/gradient passes.

    Args:
      xa: Input, branch A.
      xb: Input, branch B.
      rng: PRNG key.

    Returns:
      Mix of input branches.
    """
    if rng is None:
        rng = flax.nn.make_rng()
    gate_forward_key, gate_backward_key = jax.random.split(rng, num=2)
    gate_shape = (len(xa), 1, 1, 1)

    # Draw different interpolation factors (gate) for forward and backward pass.
    gate_forward = jax.random.uniform(
        gate_forward_key, gate_shape, dtype=jnp.float32, minval=0.0, maxval=1.0
    )
    gate_backward = jax.random.uniform(
        gate_backward_key, gate_shape, dtype=jnp.float32, minval=0.0, maxval=1.0
    )
    # Compute interpolated x for forward and backward.
    x_forward = xa * gate_forward + xb * (1.0 - gate_forward)
    x_backward = xa * gate_backward + xb * (1.0 - gate_backward)
    # Combine using stop_gradient.
    return x_backward + jax.lax.stop_gradient(x_forward - x_backward)