Gradient for constrained optimization on an L2 unit ball.

  This function projects the gradient onto the ball if you are on the boundary
  (or outside!), but leaves it untouched if you are inside the ball.

  Args:
    op: the tensorflow op we're computing the gradient for.
    grad: gradient we need to backprop

  Returns:
    (projected if necessary) gradient.

def _constrain_L2_grad(op, grad):
  """Gradient for constrained optimization on an L2 unit ball.

  This function projects the gradient onto the ball if you are on the boundary
  (or outside!), but leaves it untouched if you are inside the ball.

  Args:
    op: the tensorflow op we're computing the gradient for.
    grad: gradient we need to backprop

  Returns:
    (projected if necessary) gradient.
  """
  inp = op.inputs[0]
  inp_norm = tf.norm(inp)
  unit_inp = inp / inp_norm

  grad_projection = dot(unit_inp, grad)
  parallel_grad = unit_inp * grad_projection

  is_in_ball = tf.less_equal(inp_norm, 1)
  is_pointed_inward = tf.less(grad_projection, 0)
  allow_grad = tf.logical_or(is_in_ball, is_pointed_inward)
  clip_grad = tf.logical_not(allow_grad)

  clipped_grad = tf.cond(clip_grad, lambda: grad - parallel_grad, lambda: grad)

  return clipped_grad