code: def shakeshake2_indiv_grad(x1, x2, dy):
    
    y = shakeshake2_py(x1, x2, individual=True)
    dx = tf.gradients(ys=[y], xs=[x1, x2], grad_ys=[dy])
    return dx
docs: Overriding gradient for shake-shake of 2 tensors.