def _makeTridiagonalMatrix(self, superdiag, maindiag, subdiag):
    super_pad = [[0, 0], [0, 1], [1, 0]]
    sub_pad = [[0, 0], [1, 0], [0, 1]]

    super_part = array_ops.pad(array_ops.matrix_diag(superdiag), super_pad)
    main_part = array_ops.matrix_diag(maindiag)
    sub_part = array_ops.pad(array_ops.matrix_diag(subdiag), sub_pad)
    return super_part + main_part + sub_part

  def Test(self):
    np.random.seed(1)
    n = shape_[-1]
    batch_shape = shape_[:-2]
    np_dtype = dtype_.as_numpy_dtype
    a = np.random.uniform(
        low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
    if dtype_.is_complex:
      a += 1j * np.random.uniform(
          low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
    a += np.conj(a.T)
    a = np.tile(a, batch_shape + (1, 1))
    if dtype_ in (dtypes_lib.float32, dtypes_lib.complex64):
      atol = 1e-4
    else:
      atol = 1e-12
    np_e, np_v = np.linalg.eigh(a)
    with self.test_session():
      if compute_v_:
        tf_e, tf_v = linalg_ops.self_adjoint_eig(constant_op.constant(a))

        # Check that V*diag(E)*V^T is close to A.
        a_ev = math_ops.matmul(
            math_ops.matmul(tf_v, array_ops.matrix_diag(tf_e)),
            tf_v,
            adjoint_b=True)
        self.assertAllClose(a_ev.eval(), a, atol=atol)

        # Compare to numpy.linalg.eigh.
        CompareEigenDecompositions(self, np_e, np_v,
                                   tf_e.eval(), tf_v.eval(), atol)
      else:
        tf_e = linalg_ops.self_adjoint_eigvals(constant_op.constant(a))
        self.assertAllClose(
            np.sort(np_e, -1), np.sort(tf_e.eval(), -1), atol=atol)

 def testVector(self):
   with self.session(use_gpu=True):
     v = np.array([1.0, 2.0, 3.0])
     mat = np.diag(v)
     v_diag = array_ops.matrix_diag(v)
     self.assertEqual((3, 3), v_diag.get_shape())
     self.assertAllEqual(v_diag.eval(), mat)

  def testSampleWithBroadcastScale(self):
    # mu corresponds to a 2-batch of 3-variate normals
    mu = np.zeros([2, 3])

    # diag corresponds to no batches of 3-variate normals
    diag = np.ones([3])

    with self.test_session():
      dist = ds.VectorExponentialDiag(mu, diag, validate_args=True)

      mean = dist.mean()
      self.assertAllEqual([2, 3], mean.get_shape())
      self.assertAllClose(mu + diag, mean.eval())

      n = int(1e4)
      samps = dist.sample(n, seed=0).eval()
      samps_centered = samps - samps.mean(axis=0)
      cov_mat = array_ops.matrix_diag(diag).eval()**2
      sample_cov = np.matmul(samps_centered.transpose([1, 2, 0]),
                             samps_centered.transpose([1, 0, 2])) / n

      self.assertAllClose(mu + diag, samps.mean(axis=0),
                          atol=0.10, rtol=0.05)
      self.assertAllClose([cov_mat, cov_mat], sample_cov,
                          atol=0.10, rtol=0.05)

  def test_broadcast_matmul_and_solve(self):
    # These cannot be done in the automated (base test class) tests since they
    # test shapes that tf.matmul cannot handle.
    # In particular, tf.matmul does not broadcast.
    with self.test_session() as sess:
      x = random_ops.random_normal(shape=(2, 2, 3, 4))

      # This LinearOperatorDiag will be broadcast to (2, 2, 3, 3) during solve
      # and matmul with 'x' as the argument.
      diag = random_ops.random_uniform(shape=(2, 1, 3))
      operator = linalg.LinearOperatorDiag(diag, is_self_adjoint=True)
      self.assertAllEqual((2, 1, 3, 3), operator.shape)

      # Create a batch matrix with the broadcast shape of operator.
      diag_broadcast = array_ops.concat((diag, diag), 1)
      mat = array_ops.matrix_diag(diag_broadcast)
      self.assertAllEqual((2, 2, 3, 3), mat.get_shape())  # being pedantic.

      operator_matmul = operator.matmul(x)
      mat_matmul = math_ops.matmul(mat, x)
      self.assertAllEqual(operator_matmul.get_shape(), mat_matmul.get_shape())
      self.assertAllClose(*sess.run([operator_matmul, mat_matmul]))

      operator_solve = operator.solve(x)
      mat_solve = linalg_ops.matrix_solve(mat, x)
      self.assertAllEqual(operator_solve.get_shape(), mat_solve.get_shape())
      self.assertAllClose(*sess.run([operator_solve, mat_solve]))

 def _covariance(self):
   if (isinstance(self.scale, linalg.LinearOperatorIdentity) or
       isinstance(self.scale, linalg.LinearOperatorScaledIdentity) or
       isinstance(self.scale, linalg.LinearOperatorDiag)):
     return array_ops.matrix_diag(math_ops.square(self.scale.diag_part()))
   else:
     # TODO(b/35040238): Remove transpose once LinOp supports `transpose`.
     return self.scale.apply(array_ops.matrix_transpose(self.scale.to_dense()))

 def _testBatchVector(self, dtype):
   with self.cached_session(use_gpu=True):
     v_batch = np.array([[1.0, 0.0, 3.0], [4.0, 5.0, 6.0]]).astype(dtype)
     mat_batch = np.array([[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 3.0]],
                           [[4.0, 0.0, 0.0], [0.0, 5.0, 0.0],
                            [0.0, 0.0, 6.0]]]).astype(dtype)
     v_batch_diag = array_ops.matrix_diag(v_batch)
     self.assertEqual((2, 3, 3), v_batch_diag.get_shape())
     self.assertAllEqual(v_batch_diag.eval(), mat_batch)

def eye(
    num_rows,
    num_columns=None,
    batch_shape=None,
    dtype=dtypes.float32,
    name=None):
  """Construct an identity matrix, or a batch of matrices.

  ```python
  # Construct one identity matrix.
  tf.eye(2)
  ==> [[1., 0.],
       [0., 1.]]

  # Construct a batch of 3 identity matricies, each 2 x 2.
  # batch_identity[i, :, :] is a 2 x 2 identity matrix, i = 0, 1, 2.
  batch_identity = tf.eye(2, batch_shape=[3])

  # Construct one 2 x 3 "identity" matrix
  tf.eye(2, num_columns=3)
  ==> [[ 1.,  0.,  0.],
       [ 0.,  1.,  0.]]
  ```

  Args:
    num_rows: Non-negative `int32` scalar `Tensor` giving the number of rows
      in each batch matrix.
    num_columns: Optional non-negative `int32` scalar `Tensor` giving the number
      of columns in each batch matrix.  Defaults to `num_rows`.
    batch_shape:  `int32` `Tensor`.  If provided, returned `Tensor` will have
      leading batch dimensions of this shape.
    dtype:  The type of an element in the resulting `Tensor`
    name:  A name for this `Op`.  Defaults to "eye".

  Returns:
    A `Tensor` of shape `batch_shape + [num_rows, num_columns]`
  """
  with ops.name_scope(
      name, default_name="eye", values=[num_rows, num_columns, batch_shape]):

    batch_shape = [] if batch_shape is None else batch_shape
    batch_shape = ops.convert_to_tensor(
        batch_shape, name="shape", dtype=dtypes.int32)

    if num_columns is None:
      diag_size = num_rows
    else:
      diag_size = math_ops.minimum(num_rows, num_columns)
    diag_shape = array_ops.concat_v2((batch_shape, [diag_size]), 0)
    diag_ones = array_ops.ones(diag_shape, dtype=dtype)

    if num_columns is None:
      return array_ops.matrix_diag(diag_ones)
    else:
      shape = array_ops.concat_v2((batch_shape, [num_rows, num_columns]), 0)
      zero_matrix = array_ops.zeros(shape, dtype=dtype)
      return array_ops.matrix_set_diag(zero_matrix, diag_ones)

  def testSample(self):
    mu = [-1., 1]
    diag = [1., -2]
    with self.cached_session():
      dist = ds.MultivariateNormalDiag(mu, diag, validate_args=True)
      samps = dist.sample(int(1e3), seed=0).eval()
      cov_mat = array_ops.matrix_diag(diag).eval()**2

      self.assertAllClose(mu, samps.mean(axis=0), atol=0., rtol=0.05)
      self.assertAllClose(cov_mat, np.cov(samps.T), atol=0.05, rtol=0.05)

  def testMultivariateNormalDiagWithSoftplusStDev(self):
    mu = [-1.0, 1.0]
    diag = [-1.0, -2.0]
    with self.test_session():
      dist = distributions.MultivariateNormalDiagWithSoftplusStDev(mu, diag)
      samps = dist.sample(1000, seed=0).eval()
      cov_mat = array_ops.matrix_diag(nn_ops.softplus(diag)).eval()**2

      self.assertAllClose(mu, samps.mean(axis=0), atol=0.1)
      self.assertAllClose(cov_mat, np.cov(samps.T), atol=0.1)

 def _covariance(self):
   # Let
   #   W = (w1,...,wk), with wj ~ iid Exponential(0, 1).
   # Then this distribution is
   #   X = loc + LW,
   # and then since Cov(wi, wj) = 1 if i=j, and 0 otherwise,
   #   Cov(X) = L Cov(W W^T) L^T = L L^T.
   if distribution_util.is_diagonal_scale(self.scale):
     return array_ops.matrix_diag(math_ops.square(self.scale.diag_part()))
   else:
     return self.scale.matmul(self.scale.to_dense(), adjoint_arg=True)

 def testGrad(self):
   shapes = ((3,), (7, 4))
   with self.session(use_gpu=True):
     for shape in shapes:
       x = constant_op.constant(np.random.rand(*shape), np.float32)
       y = array_ops.matrix_diag(x)
       error = gradient_checker.compute_gradient_error(x,
                                                       x.get_shape().as_list(),
                                                       y,
                                                       y.get_shape().as_list())
       self.assertLess(error, 1e-4)

  def testSample(self):
    mu = [-1., 1]
    diag = [1., -2]
    with self.test_session():
      dist = ds.VectorLaplaceDiag(mu, diag, validate_args=True)
      samps = dist.sample(int(1e4), seed=0).eval()
      cov_mat = 2. * array_ops.matrix_diag(diag).eval()**2

      self.assertAllClose(mu, samps.mean(axis=0),
                          atol=0., rtol=0.05)
      self.assertAllClose(cov_mat, np.cov(samps.T),
                          atol=0.05, rtol=0.05)

 def _build_operator_and_mat(self, batch_shape, k, dtype=np.float64):
   # Build an identity matrix with right shape and dtype.
   # Build an operator that should act the same way.
   batch_shape = list(batch_shape)
   diag_shape = batch_shape + [k]
   matrix_shape = batch_shape + [k, k]
   diag = array_ops.ones(diag_shape, dtype=dtype)
   scale = constant_op.constant(2.0, dtype=dtype)
   scaled_identity_matrix = scale * array_ops.matrix_diag(diag)
   operator = operator_pd_identity.OperatorPDIdentity(
       matrix_shape, dtype, scale=scale)
   return operator, scaled_identity_matrix.eval()

  def testSample(self):
    mu = [-2., 1]
    diag = [1., -2]
    with self.cached_session():
      dist = ds.VectorExponentialDiag(mu, diag, validate_args=True)
      samps = dist.sample(int(1e4), seed=0).eval()
      cov_mat = array_ops.matrix_diag(diag).eval()**2

      self.assertAllClose([-2 + 1, 1. - 2], samps.mean(axis=0),
                          atol=0., rtol=0.05)
      self.assertAllClose(cov_mat, np.cov(samps.T),
                          atol=0.05, rtol=0.05)

def _SelfAdjointEigV2Grad(op, grad_e, grad_v):
  """Gradient for SelfAdjointEigV2."""
  e = op.outputs[0]
  compute_v = op.get_attr("compute_v")
  # a = op.inputs[0], which satisfies
  # a[...,:,:] * v[...,:,i] = e[...,i] * v[...,i]
  with ops.control_dependencies([grad_e, grad_v]):
    if compute_v:
      v = op.outputs[1]
      # Construct the matrix f(i,j) = (i != j ? 1 / (e_i - e_j) : 0).
      # Notice that because of the term involving f, the gradient becomes
      # infinite (or NaN in practice) when eigenvalues are not unique.
      # Mathematically this should not be surprising, since for (k-fold)
      # degenerate eigenvalues, the corresponding eigenvectors are only defined
      # up to arbitrary rotation in a (k-dimensional) subspace.
      f = array_ops.matrix_set_diag(
          math_ops.reciprocal(
              array_ops.expand_dims(e, -2) - array_ops.expand_dims(e, -1)),
          array_ops.zeros_like(e))
      grad_a = math_ops.matmul(
          v,
          math_ops.matmul(
              array_ops.matrix_diag(grad_e) +
              f * math_ops.matmul(v, grad_v, adjoint_a=True),
              v,
              adjoint_b=True))
    else:
      _, v = linalg_ops.self_adjoint_eig(op.inputs[0])
      grad_a = math_ops.matmul(v,
                               math_ops.matmul(
                                   array_ops.matrix_diag(grad_e),
                                   v,
                                   adjoint_b=True))
    # The forward op only depends on the lower triangular part of a, so here we
    # symmetrize and take the lower triangle
    grad_a = array_ops.matrix_band_part(
        grad_a + math_ops.conj(array_ops.matrix_transpose(grad_a)), -1, 0)
    grad_a = array_ops.matrix_set_diag(grad_a,
                                       0.5 * array_ops.matrix_diag_part(grad_a))
    return grad_a

 def _covariance(self):
   # Let
   #   W = (w1,...,wk), with wj ~ iid Laplace(0, 1).
   # Then this distribution is
   #   X = loc + LW,
   # and since E[X] = loc,
   #   Cov(X) = E[LW W^T L^T] = L E[W W^T] L^T.
   # Since E[wi wj] = 0 if i != j, and 2 if i == j, we have
   #   Cov(X) = 2 LL^T
   if distribution_util.is_diagonal_scale(self.scale):
     return 2. * array_ops.matrix_diag(math_ops.square(self.scale.diag_part()))
   else:
     return 2. * self.scale.matmul(self.scale.to_dense(), adjoint_arg=True)

 def testBatchVector(self):
   with self.test_session(use_gpu=self._use_gpu):
     v_batch = np.array([[1.0, 2.0, 3.0],
                         [4.0, 5.0, 6.0]])
     mat_batch = np.array(
         [[[1.0, 0.0, 0.0],
           [0.0, 2.0, 0.0],
           [0.0, 0.0, 3.0]],
          [[4.0, 0.0, 0.0],
           [0.0, 5.0, 0.0],
           [0.0, 0.0, 6.0]]])
     v_batch_diag = array_ops.matrix_diag(v_batch)
     self.assertEqual((2, 3, 3), v_batch_diag.get_shape())
     self.assertAllEqual(v_batch_diag.eval(), mat_batch)

  def _operator_and_matrix(self, build_info, dtype, use_placeholder):
    shape = list(build_info.shape)
    diag = linear_operator_test_util.random_sign_uniform(
        shape[:-1], minval=1., maxval=2., dtype=dtype)

    lin_op_diag = diag

    if use_placeholder:
      lin_op_diag = array_ops.placeholder_with_default(diag, shape=None)

    operator = linalg.LinearOperatorDiag(lin_op_diag)

    matrix = array_ops.matrix_diag(diag)

    return operator, matrix

def eye(num_rows,
        num_columns=None,
        batch_shape=None,
        dtype=dtypes.float32,
        name=None):
  """Construct an identity matrix, or a batch of matrices.

  See `linalg_ops.eye`.
  """
  with ops.name_scope(
      name, default_name='eye', values=[num_rows, num_columns, batch_shape]):
    is_square = num_columns is None
    batch_shape = [] if batch_shape is None else batch_shape
    num_columns = num_rows if num_columns is None else num_columns

    # We cannot statically infer what the diagonal size should be:
    if (isinstance(num_rows, ops.Tensor) or
        isinstance(num_columns, ops.Tensor)):
      diag_size = math_ops.minimum(num_rows, num_columns)
    else:
      # We can statically infer the diagonal size, and whether it is square.
      if not isinstance(num_rows, compat.integral_types) or not isinstance(
          num_columns, compat.integral_types):
        raise TypeError(
            'num_rows and num_columns must be positive integer values.')
      is_square = num_rows == num_columns
      diag_size = np.minimum(num_rows, num_columns)

    # We can not statically infer the shape of the tensor.
    if isinstance(batch_shape, ops.Tensor) or isinstance(diag_size, ops.Tensor):
      batch_shape = ops.convert_to_tensor(
          batch_shape, name='shape', dtype=dtypes.int32)
      diag_shape = array_ops.concat((batch_shape, [diag_size]), axis=0)
      if not is_square:
        shape = array_ops.concat((batch_shape, [num_rows, num_columns]), axis=0)
    # We can statically infer everything.
    else:
      batch_shape = list(batch_shape)
      diag_shape = batch_shape + [diag_size]
      if not is_square:
        shape = batch_shape + [num_rows, num_columns]

    diag_ones = array_ops.ones(diag_shape, dtype=dtype)
    if is_square:
      return array_ops.matrix_diag(diag_ones)
    else:
      zero_matrix = array_ops.zeros(shape, dtype=dtype)
      return array_ops.matrix_set_diag(zero_matrix, diag_ones)