def _ZetaGrad(op, grad):
    """Returns gradient of zeta(x, q) with respect to x and q."""
    # TODO(tillahoffmann): Add derivative with respect to x
    x = op.inputs[0]
    q = op.inputs[1]
    # Broadcast gradients
    sx = array_ops.shape(x)
    sq = array_ops.shape(q)
    unused_rx, rq = gen_array_ops._broadcast_gradient_args(sx, sq)
    # Evaluate gradient
    with ops.control_dependencies([grad.op]):
        partial_q = -x * math_ops.zeta(x + 1, q)
        return (None, array_ops.reshape(math_ops.reduce_sum(partial_q * grad, rq), sq))

def _ZetaGrad(op, grad):
  """Returns gradient of zeta(x, q) with respect to x and q."""
  # TODO(tillahoffmann): Add derivative with respect to x
  x = op.inputs[0]
  q = op.inputs[1]
  # Broadcast gradients
  sx = array_ops.shape(x)
  sq = array_ops.shape(q)
  # pylint: disable=protected-access
  unused_rx, rq = gen_array_ops._broadcast_gradient_args(sx, sq)
  # pylint: enable=protected-access
  # Evaluate gradient
  with ops.control_dependencies([grad]):
    x = math_ops.conj(x)
    q = math_ops.conj(q)
    partial_q = -x * math_ops.zeta(x + 1, q)
    # TODO(b/36815900): Mark None return values as NotImplemented
    return (None,
            array_ops.reshape(math_ops.reduce_sum(partial_q * grad, rq), sq))

 def safe_zeta(x, y):
   return math_ops.zeta(x * x + 1, y * y)