def _compareRealImag(self, cplx, use_gpu):
   np_real, np_imag = np.real(cplx), np.imag(cplx)
   np_zeros = np_real * 0
   with self.test_session(use_gpu=use_gpu,
                          force_gpu=use_gpu and test_util.is_gpu_available()):
     inx = ops.convert_to_tensor(cplx)
     tf_real = math_ops.real(inx)
     tf_imag = math_ops.imag(inx)
     tf_real_real = math_ops.real(tf_real)
     tf_imag_real = math_ops.imag(tf_real)
     self.assertAllEqual(np_real, self.evaluate(tf_real))
     self.assertAllEqual(np_imag, self.evaluate(tf_imag))
     self.assertAllEqual(np_real, self.evaluate(tf_real_real))
     self.assertAllEqual(np_zeros, self.evaluate(tf_imag_real))

  def _compareRealImag(self, cplx, use_gpu):
    np_real, np_imag = np.real(cplx), np.imag(cplx)
    np_zeros = np_real * 0

    with test_util.device(use_gpu=use_gpu):
      inx = ops.convert_to_tensor(cplx)
      tf_real = math_ops.real(inx)
      tf_imag = math_ops.imag(inx)
      tf_real_real = math_ops.real(tf_real)
      tf_imag_real = math_ops.imag(tf_real)
      self.assertAllEqual(np_real, self.evaluate(tf_real))
      self.assertAllEqual(np_imag, self.evaluate(tf_imag))
      self.assertAllEqual(np_real, self.evaluate(tf_real_real))
      self.assertAllEqual(np_zeros, self.evaluate(tf_imag_real))

  def test_defining_spd_operator_by_taking_real_part(self):
    with self.cached_session() as sess:
      # S is real and positive.
      s = linear_operator_test_util.random_uniform(
          shape=(10, 2, 3, 4), dtype=dtypes.float32, minval=1., maxval=2.)

      # Let S = S1 + S2, the Hermitian and anti-hermitian parts.
      # S1 = 0.5 * (S + S^H), S2 = 0.5 * (S - S^H),
      # where ^H is the Hermitian transpose of the function:
      #    f(n0, n1, n2)^H := ComplexConjugate[f(N0-n0, N1-n1, N2-n2)].
      # We want to isolate S1, since
      #   S1 is Hermitian by construction
      #   S1 is real since S is
      #   S1 is positive since it is the sum of two positive kernels

      # IDFT[S] = IDFT[S1] + IDFT[S2]
      #         =      H1  +      H2
      # where H1 is real since it is Hermitian,
      # and H2 is imaginary since it is anti-Hermitian.
      ifft_s = fft_ops.ifft3d(math_ops.cast(s, dtypes.complex64))

      # Throw away H2, keep H1.
      real_ifft_s = math_ops.real(ifft_s)

      # This is the perfect spectrum!
      # spectrum = DFT[H1]
      #          = S1,
      fft_real_ifft_s = fft_ops.fft3d(
          math_ops.cast(real_ifft_s, dtypes.complex64))

      # S1 is Hermitian ==> operator is real.
      # S1 is real ==> operator is self-adjoint.
      # S1 is positive ==> operator is positive-definite.
      operator = linalg.LinearOperatorCirculant3D(fft_real_ifft_s)

      # Allow for complex output so we can check operator has zero imag part.
      self.assertEqual(operator.dtype, dtypes.complex64)
      matrix, matrix_t = sess.run([
          operator.to_dense(),
          array_ops.matrix_transpose(operator.to_dense())
      ])
      operator.assert_positive_definite().run()  # Should not fail.
      np.testing.assert_allclose(0, np.imag(matrix), atol=1e-6)
      self.assertAllClose(matrix, matrix_t)

      # Just to test the theory, get S2 as well.
      # This should create an imaginary operator.
      # S2 is anti-Hermitian ==> operator is imaginary.
      # S2 is real ==> operator is self-adjoint.
      imag_ifft_s = math_ops.imag(ifft_s)
      fft_imag_ifft_s = fft_ops.fft3d(
          1j * math_ops.cast(imag_ifft_s, dtypes.complex64))
      operator_imag = linalg.LinearOperatorCirculant3D(fft_imag_ifft_s)

      matrix, matrix_h = sess.run([
          operator_imag.to_dense(),
          array_ops.matrix_transpose(math_ops.conj(operator_imag.to_dense()))
      ])
      self.assertAllClose(matrix, matrix_h)
      np.testing.assert_allclose(0, np.real(matrix), atol=1e-7)

  def _compareMulGradient(self, data):
    # data is a float matrix of shape [n, 4].  data[:, 0], data[:, 1],
    # data[:, 2], data[:, 3] are real parts of x, imaginary parts of
    # x, real parts of y and imaginary parts of y.
    with self.cached_session():
      inp = ops.convert_to_tensor(data)
      xr, xi, yr, yi = array_ops.split(value=inp, num_or_size_splits=4, axis=1)

      def vec(x):  # Reshape to a vector
        return array_ops.reshape(x, [-1])

      xr, xi, yr, yi = vec(xr), vec(xi), vec(yr), vec(yi)

      def cplx(r, i):  # Combine to a complex vector
        return math_ops.complex(r, i)

      x, y = cplx(xr, xi), cplx(yr, yi)
      # z is x times y in complex plane.
      z = x * y
      # Defines the loss function as the sum of all coefficients of z.
      loss = math_ops.reduce_sum(math_ops.real(z) + math_ops.imag(z))
      epsilon = 0.005
      jacob_t, jacob_n = gradient_checker.compute_gradient(
          inp, list(data.shape), loss, [1], x_init_value=data, delta=epsilon)
    self.assertAllClose(jacob_t, jacob_n, rtol=epsilon, atol=epsilon)

  def _trace(self):
    # The diagonal of the [[nested] block] circulant operator is the mean of
    # the spectrum.
    # Proof:  For the [0,...,0] element, this follows from the IDFT formula.
    # Then the result follows since all diagonal elements are the same.

    # Therefore, the trace is the sum of the spectrum.

    # Get shape of diag along with the axis over which to reduce the spectrum.
    # We will reduce the spectrum over all block indices.
    if self.spectrum.get_shape().is_fully_defined():
      spec_rank = self.spectrum.get_shape().ndims
      axis = np.arange(spec_rank - self.block_depth, spec_rank, dtype=np.int32)
    else:
      spec_rank = array_ops.rank(self.spectrum)
      axis = math_ops.range(spec_rank - self.block_depth, spec_rank)

    # Real diag part "re_d".
    # Suppose spectrum.shape = [B1,...,Bb, N1, N2]
    # self.shape = [B1,...,Bb, N, N], with N1 * N2 = N.
    # re_d_value.shape = [B1,...,Bb]
    re_d_value = math_ops.reduce_sum(math_ops.real(self.spectrum), axis=axis)

    if not self.dtype.is_complex:
      return math_ops.cast(re_d_value, self.dtype)

    # Imaginary part, "im_d".
    if self.is_self_adjoint:
      im_d_value = 0.
    else:
      im_d_value = math_ops.reduce_sum(math_ops.imag(self.spectrum), axis=axis)

    return math_ops.cast(math_ops.complex(re_d_value, im_d_value), self.dtype)

def _ComplexGrad(op, grad):
  """Returns the real and imaginary components of 'grad', respectively."""
  x = op.inputs[0]
  y = op.inputs[1]
  sx = array_ops.shape(x)
  sy = array_ops.shape(y)
  rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy)
  return (array_ops.reshape(math_ops.reduce_sum(math_ops.real(grad), rx), sx),
          array_ops.reshape(math_ops.reduce_sum(math_ops.imag(grad), ry), sy))

 def test_simple_hermitian_spectrum_gives_operator_with_zero_imag_part(self):
   with self.test_session():
     spectrum = math_ops.cast([1., 1j, -1j], dtypes.complex64)
     operator = linalg.LinearOperatorCirculant(
         spectrum, input_output_dtype=dtypes.complex64)
     matrix = operator.to_dense()
     imag_matrix = math_ops.imag(matrix)
     eps = np.finfo(np.float32).eps
     np.testing.assert_allclose(0, imag_matrix.eval(), rtol=0, atol=eps * 3)

def _AngleGrad(op, grad):
  """Returns -grad / (Im(x) + iRe(x))"""
  x = op.inputs[0]
  with ops.control_dependencies([grad]):
    re = math_ops.real(x)
    im = math_ops.imag(x)
    z = math_ops.reciprocal(math_ops.complex(im, re))
    zero = constant_op.constant(0, dtype=grad.dtype)
    complex_grad = math_ops.complex(grad, zero)
    return -complex_grad * z

 def test_hermitian_spectrum_gives_operator_with_zero_imag_part(self):
   with self.cached_session():
     # Make spectrum the FFT of a real convolution kernel h.  This ensures that
     # spectrum is Hermitian.
     h = linear_operator_test_util.random_normal(shape=(3, 4))
     spectrum = fft_ops.fft(math_ops.cast(h, dtypes.complex64))
     operator = linalg.LinearOperatorCirculant(
         spectrum, input_output_dtype=dtypes.complex64)
     matrix = operator.to_dense()
     imag_matrix = math_ops.imag(matrix)
     eps = np.finfo(np.float32).eps
     np.testing.assert_allclose(
         0, self.evaluate(imag_matrix), rtol=0, atol=eps * 3 * 4)

def _assert_imag_part_zero(x, message=None):
  """Assert that floating or complex 'x' is real."""
  dtype = x.dtype.base_dtype
  if dtype.is_floating:
    return control_flow_ops.no_op()

  if not dtype.is_complex:
    raise TypeError(
        "imag_part_zero only handles float or complex types.  Found: %s"
        % dtype)

  zero = ops.convert_to_tensor(0, dtype=dtype.real_dtype)
  return check_ops.assert_equal(zero, math_ops.imag(x), message=message)

  def test_real_hermitian_spectrum_gives_real_symmetric_operator(self):
    with self.cached_session() as sess:
      # This is a real and hermitian spectrum.
      spectrum = [[1., 2., 2.], [3., 4., 4.], [3., 4., 4.]]
      operator = linalg.LinearOperatorCirculant(spectrum)

      matrix_tensor = operator.to_dense()
      self.assertEqual(matrix_tensor.dtype, dtypes.complex64)
      matrix_t = array_ops.matrix_transpose(matrix_tensor)
      imag_matrix = math_ops.imag(matrix_tensor)
      matrix, matrix_transpose, imag_matrix = sess.run(
          [matrix_tensor, matrix_t, imag_matrix])

      np.testing.assert_allclose(0, imag_matrix, atol=1e-6)
      self.assertAllClose(matrix, matrix_transpose, atol=0)

 def _compareGradient(self, x):
   # x[:, 0] is real, x[:, 1] is imag.  We combine real and imag into
   # complex numbers. Then, we extract real and imag parts and
   # computes the squared sum. This is obviously the same as sum(real
   # * real) + sum(imag * imag). We just want to make sure the
   # gradient function is checked.
   with self.cached_session():
     inx = ops.convert_to_tensor(x)
     real, imag = array_ops.split(value=inx, num_or_size_splits=2, axis=1)
     real, imag = array_ops.reshape(real, [-1]), array_ops.reshape(imag, [-1])
     cplx = math_ops.complex(real, imag)
     cplx = math_ops.conj(cplx)
     loss = math_ops.reduce_sum(math_ops.square(
         math_ops.real(cplx))) + math_ops.reduce_sum(
             math_ops.square(math_ops.imag(cplx)))
     epsilon = 1e-3
     jacob_t, jacob_n = gradient_checker.compute_gradient(
         inx, list(x.shape), loss, [1], x_init_value=x, delta=epsilon)
   self.assertAllClose(jacob_t, jacob_n, rtol=epsilon, atol=epsilon)

def assert_zero_imag_part(x, message=None, name="assert_zero_imag_part"):
  """Returns `Op` that asserts Tensor `x` has no non-zero imaginary parts.

  Args:
    x:  Numeric `Tensor`, real, integer, or complex.
    message:  A string message to prepend to failure message.
    name:  A name to give this `Op`.

  Returns:
    An `Op` that asserts `x` has no entries with modulus zero.
  """
  with ops.name_scope(name, values=[x]):
    x = ops.convert_to_tensor(x, name="x")
    dtype = x.dtype.base_dtype

    if dtype.is_floating:
      return control_flow_ops.no_op()

    zero = ops.convert_to_tensor(0, dtype=dtype.real_dtype)
    return check_ops.assert_equal(zero, math_ops.imag(x), message=message)

  def assert_hermitian_spectrum(self, name="assert_hermitian_spectrum"):
    """Returns an `Op` that asserts this operator has Hermitian spectrum.

    This operator corresponds to a real-valued matrix if and only if its
    spectrum is Hermitian.

    Args:
      name:  A name to give this `Op`.

    Returns:
      An `Op` that asserts this operator has Hermitian spectrum.
    """
    eps = np.finfo(self.dtype.real_dtype.as_numpy_dtype).eps
    with self._name_scope(name):
      # Assume linear accumulation of error.
      max_err = eps * self.domain_dimension_tensor()
      imag_convolution_kernel = math_ops.imag(self.convolution_kernel())
      return check_ops.assert_less(
          math_ops.abs(imag_convolution_kernel),
          max_err,
          message="Spectrum was not Hermitian")

def _abs_square(x):
  if x.dtype.is_complex:
    return math_ops.square(math_ops.real(x)) + math_ops.square(math_ops.imag(x))
  else:
    return math_ops.square(x)

def _ComplexGrad(_, grad):
  """Returns the real and imaginary components of 'grad', respectively."""
  return math_ops.real(grad), math_ops.imag(grad)

 def _assert_self_adjoint(self):
   imag_multiplier = math_ops.imag(self.multiplier)
   return check_ops.assert_equal(
       array_ops.zeros_like(imag_multiplier),
       imag_multiplier,
       message="LinearOperator was not self-adjoint")