message: Add missing diagonal term
diff: @@ -642,8 +642,9 @@ def _SafeReciprocal(x, epsilon=1E-20):
  @ops.RegisterGradient("Eig")
  def _EigGrad(op, grad_e, grad_v):
  """Gradient for Eig.
  - Based on report by Mike Giles
  - https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
  + Based on eq. 4.77 from paper by
  + Christoph Boeddeker et al.
  + https://arxiv.org/abs/1701.00392
  See also
  "Computation of eigenvalue and eigenvector derivatives
  for a general complex-valued eigensystem" by Nico van der Aa.
  @@ -667,13 +668,16 @@ def _EigGrad(op, grad_e, grad_v):
  _SafeReciprocal(
  array_ops.expand_dims(e, -2) - array_ops.expand_dims(e, -1)),
  array_ops.zeros_like(e))
  - grad_a = math_ops.matmul(
  - w,
  - math_ops.matmul(
  - array_ops.matrix_diag(grad_e) +
  - f * math_ops.matmul(v, grad_v, adjoint_a=True),
  - v,
  - adjoint_b=True))
  +
  + vgv = math_ops.matmul(v, grad_v, adjoint_a=True)
  + mid = array_ops.matrix_diag(grad_e)
  + diag_grad_part = array_ops.matrix_set_diag(
  + array_ops.zeros_like(vgv),
  + array_ops.matrix_diag_part(
  + math_ops.cast(math_ops.real(vgv), vgv.dtype)))
  + mid += f * (vgv - math_ops.matmul(math_ops.matmul(v, v, adjoint_a=True),
  + diag_grad_part))
  + grad_a = math_ops.matmul(w, math_ops.matmul(mid, v, adjoint_b=True))
  else:
  _, v = linalg_ops.self_adjoint_eig(op.inputs[0])
  w = _linalg.adjoint(linalg_ops.matrix_inverse(v))
  