def testWrongDimensions(self):
     # The matrix and right-hand sides should have the same number of rows.
     with self.test_session():
         matrix = tf.constant([[1.0, 0.0], [0.0, 1.0]])
         rhs = tf.constant([[1.0, 0.0]])
         with self.assertRaises(ValueError):
             tf.matrix_solve(matrix, rhs)

 def testNotInvertible(self):
     # The input should be invertible.
     with self.test_session():
         with self.assertRaisesOpError("Input matrix is not invertible."):
             # All rows of the matrix below add to zero
             matrix = tf.constant([[1.0, 0.0, -1.0], [-1.0, 1.0, 0.0], [0.0, -1.0, 1.0]])
             tf.matrix_solve(matrix, matrix).eval()

 def testNonSquareMatrix(self):
     # When the solve of a non-square matrix is attempted we should return
     # an error
     with self.test_session():
         with self.assertRaises(ValueError):
             matrix = tf.constant([[1.0, 2.0, 3.0], [3.0, 4.0, 5.0]])
             tf.matrix_solve(matrix, matrix)

def pred(X,X_m_1,mu,len_sc_1,noise_1):
        Kmm=h.tf_SE_K(X_m_1,X_m_1,len_sc_1,noise_1)
        Knm=h.tf_SE_K(X,X_m_1,len_sc_1,noise_1)
        posterior_mean= h.Mul(Knm,tf.matrix_solve(Kmm,mu))
        K_nn=h.tf_SE_K(X,X,len_sc_1,noise_1)
        full_cov=K_nn-h.Mul(Knm,tf.matrix_solve(Kmm,tf.transpose(Knm)))
        posterior_cov=tf.diag_part(full_cov)
        return posterior_mean,tf.reshape(posterior_cov,[N,1]),full_cov

def predict(K_mn,sigma,K_mm,K_nn):
    # predicitions
    K_nm=tf.transpose(K_mn)
    Sig_Inv=1e-1*np.eye(M)+K_mm+K_mnnm_2/tf.square(sigma)
    mu_post=h.Mul(tf.matrix_solve(Sig_Inv,K_mn),Ytr)/tf.square(sigma)
    mean=h.Mul(K_nm,mu_post)
    variance=K_nn-h.Mul(K_nm,h.safe_chol(K_mm,K_mn))+h.Mul(K_nm,tf.matrix_solve(Sig_Inv,K_mn))
    var_terms=2*tf.sqrt(tf.reshape(tf.diag_part(variance)+tf.square(sigma),[N,1]))

    return mean, var_terms

    def __init__(self, pars, vars, eqns):
        self.pars = pars
        self.vars = vars
        self.eqns = eqns

        # size
        self.par_sz = {par: int(par.get_shape()[0]) for par in pars}
        self.var_sz = {var: int(var.get_shape()[0]) for var in vars}
        self.eqn_sz = {eqn: int(eqn.get_shape()[0]) for eqn in eqns}

        # equation system
        self.parvec = tf.concat(pars, 0)
        self.varvec = tf.concat(vars, 0)
        self.eqnvec = tf.concat(eqns, 0)
        self.error = tf.reduce_max(tf.abs(self.eqnvec))

        # gradients
        self.parjac = tf.concat([tf.concat([jacobian(eqn, x) for x in pars], 1) for eqn in eqns], 0)
        self.varjac = tf.concat([tf.concat([jacobian(eqn, x) for x in vars], 1) for eqn in eqns], 0)

        # newton steps
        self.newton_step = -tf.squeeze(tf.matrix_solve(self.varjac, tf.expand_dims(self.eqnvec, 1)))
        self.newton_dvars = tf.split(self.newton_step, list(self.var_sz.values()), 0)
        self.newton_update = [tf.assign(v, v+s) for v, s in zip(self.vars, self.newton_dvars)]

        # homotopy
        self.tv = tf.placeholder(dtype=tf.float64)
        self.par0 = [tf.Variable(np.zeros(p.shape)) for p in pars]
        self.par1 = [tf.Variable(np.zeros(p.shape)) for p in pars]

        # path gen
        self.path_assign = tf.group(*[p.assign((1-self.tv)*p0 + self.tv*p1)
            for p, p0, p1 in zip(pars, self.par0, self.par1)])

  def _verifySolve(self, x, y, batch_dims=None):
    for adjoint in False, True:
      for np_type in [np.float32, np.float64]:
        a = x.astype(np_type)
        b = y.astype(np_type)
        if adjoint:
          a_np = np.conj(np.transpose(a))
        else:
          a_np = a
        if batch_dims is not None:
          a = np.tile(a, batch_dims + [1, 1])
          a_np = np.tile(a_np, batch_dims + [1, 1])
          b = np.tile(b, batch_dims + [1, 1])

        np_ans = np.linalg.solve(a_np, b)
        with self.test_session():
          # Test the batch version, which works for ndim >= 2
          tf_ans = tf.batch_matrix_solve(a, b, adjoint=adjoint)
          out = tf_ans.eval()
          self.assertEqual(tf_ans.get_shape(), out.shape)
          self.assertEqual(np_ans.shape, out.shape)
          self.assertAllClose(np_ans, out)

          if a.ndim == 2:
            # Test the simple version
            tf_ans = tf.matrix_solve(a, b, adjoint=adjoint)
            out = tf_ans.eval()
            self.assertEqual(out.shape, tf_ans.get_shape())
            self.assertEqual(np_ans.shape, out.shape)
            self.assertAllClose(np_ans, out)

    def __init__(self, pars, vars, eqns):
        self.pars = pars
        self.vars = vars
        self.eqns = eqns

        # size
        self.par_sz = {par: int(par.get_shape()[0]) for par in pars}
        self.var_sz = {var: int(var.get_shape()[0]) for var in vars}
        self.eqn_sz = {eqn: int(eqn.get_shape()[0]) for eqn in eqns}

        # equation system
        self.parvec = tf.concat(pars, 0)
        self.varvec = tf.concat(vars, 0)
        self.eqnvec = tf.concat(eqns, 0)
        self.error = tf.reduce_max(tf.abs(self.eqnvec))

        # gradients
        self.parjac = tf.concat([tf.concat([jacobian(eqn, x) for x in pars], 1) for eqn in eqns], 0)
        self.varjac = tf.concat([tf.concat([jacobian(eqn, x) for x in vars], 1) for eqn in eqns], 0)

        # newton steps
        self.newton_step = -tf.squeeze(tf.matrix_solve(self.varjac, tf.expand_dims(self.eqnvec, 1)))
        self.newton_dvars = tf.split(self.newton_step, list(self.var_sz.values()), 0)
        self.newton_update = [tf.assign(v, v+s) for v, s in zip(self.vars, self.newton_dvars)]

        # target param
        self.tpars = [tf.zeros_like(p) for p in pars]
        self.tparvec = tf.concat(self.tpars, 0)

  def test_broadcast_apply_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 = tf.random_normal(shape=(2, 2, 3, 4))

      # This LinearOperatorDiag will be brodacast to (2, 2, 3, 3) during solve
      # and apply with 'x' as the argument.
      diag = tf.random_uniform(shape=(2, 1, 3))
      operator = linalg.LinearOperatorDiag(diag)
      self.assertAllEqual((2, 1, 3, 3), operator.shape)

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

      operator_apply = operator.apply(x)
      mat_apply = tf.matmul(mat, x)
      self.assertAllEqual(operator_apply.get_shape(), mat_apply.get_shape())
      self.assertAllClose(*sess.run([operator_apply, mat_apply]))

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

  def _verifySolve(self, x, y, batch_dims=None):
    for adjoint in False, True:
      for np_type in [np.float32, np.float64, np.complex64, np.complex128]:
        if np_type is [np.float32, np.float64]:
          a = x.real().astype(np_type)
          b = y.real().astype(np_type)
        else:
          a = x.astype(np_type)
          b = y.astype(np_type)
        if adjoint:
          a_np = np.conj(np.transpose(a))
        else:
          a_np = a
        if batch_dims is not None:
          a = np.tile(a, batch_dims + [1, 1])
          a_np = np.tile(a_np, batch_dims + [1, 1])
          b = np.tile(b, batch_dims + [1, 1])

        np_ans = np.linalg.solve(a_np, b)
        with self.test_session():
          tf_ans = tf.matrix_solve(a, b, adjoint=adjoint)
          out = tf_ans.eval()
          self.assertEqual(tf_ans.get_shape(), out.shape)
          self.assertEqual(np_ans.shape, out.shape)
          self.assertAllClose(np_ans, out)

 def testBatchResultSize(self):
   # 3x3x3 matrices, 3x3x1 right-hand sides.
   matrix = np.array([1., 2., 3., 4., 5., 6., 7., 8., 9.] * 3).reshape(3, 3, 3)
   rhs = np.array([1., 2., 3.] * 3).reshape(3, 3, 1)
   answer = tf.matrix_solve(matrix, rhs)
   ls_answer = tf.matrix_solve_ls(matrix, rhs)
   self.assertEqual(ls_answer.get_shape(), [3, 3, 1])
   self.assertEqual(answer.get_shape(), [3, 3, 1])

def predict2():
    # predicitions
    cov=h.Mul(K_mm_2,tf.matrix_inverse(K_mm_2+K_mnnm_2/tf.square(sigma_2)),K_mm_2)
    cov_chol=tf.cholesky(cov)
    mu=h.Mul(K_mm_2,tf.cholesky_solve(cov_chol,K_mn_2),Ytr)/tf.square(sigma_2)
    mean=h.Mul(K_nm_2,tf.matrix_solve(K_mm_1,mu))
    variance=K_nn_2-h.Mul(K_nm_2,h.safe_chol(K_mm_2,tf.transpose(K_nm_2)))
    var_terms=2*tf.sqrt(tf.reshape(tf.diag_part(variance)+tf.square(sigma_2),[N,1]))
    return mean, var_terms

    def _logp(self, result, prior_mean, prior_cov,
                 transition_mat, transition_mean, transition_cov,
                 observation_mat=None, observation_mean=None, observation_cov=None):
    
        # define the Kalman filtering calculation within the TF graph
        if observation_mean is not None:
            observation_mean = tf.reshape(observation_mean, (self.K, 1))
            
        transition_mean = tf.reshape(transition_mean, (self.D, 1))
        
        pred_mean = tf.reshape(prior_mean, (self.D, 1))
        pred_cov = prior_cov

        filtered_means = []
        filtered_covs = []
        step_logps = []
        observations = tf.unpack(result)
        for t in range(self.T):
            obs_t = tf.reshape(observations[t], (self.K, 1))

            if observation_mat is not None:

                tmp = tf.matmul(observation_mat, pred_cov)
                S = tf.matmul(tmp, tf.transpose(observation_mat)) + observation_cov
                # TODO optimize this to not use an explicit matrix inverse
                #Sinv = tf.matrix_inverse(S)
                #gain = tf.matmul(pred_cov, tf.matmul(tf.transpose(observation_mat), Sinv))

                # todo worth implementing cholsolve explicitly?
                gain = tf.matmul(pred_cov, tf.transpose(tf.matrix_solve(S, observation_mat)))
                
                y = obs_t - tf.matmul(observation_mat, pred_mean) - observation_mean
                updated_mean = pred_mean + tf.matmul(gain, y)
                updated_cov = pred_cov - tf.matmul(gain, tmp)
            else:
                updated_mean = obs_t
                updated_cov = tf.zeros_like(pred_cov)
                S = pred_cov
                y = obs_t - pred_mean
                
            step_logp = bf.dists.multivariate_gaussian_log_density(y, 0, S)

            filtered_means.append(updated_mean)
            filtered_covs.append(updated_cov)
            step_logps.append(step_logp)

            if t < self.T-1:
                pred_mean = tf.matmul(transition_mat, updated_mean) + transition_mean
                pred_cov = tf.matmul(transition_mat, tf.matmul(updated_cov, tf.transpose(transition_mat))) + transition_cov

        self.filtered_means = filtered_means
        self.filtered_covs = filtered_covs
        self.step_logps = tf.pack(step_logps)
        logp = tf.reduce_sum(self.step_logps)

        return logp

    def testSolve(self):
        with self.test_session():
            for batch_shape in [(), (2, 3)]:
                for k in [1, 4]:
                    operator, mat = self._build_operator_and_mat(batch_shape, k)

                    # Work with 5 simultaneous systems.  5 is arbitrary.
                    x = self._rng.randn(*(batch_shape + (k, 5)))

                    self._compare_results(expected=tf.matrix_solve(mat, x).eval(), actual=operator.solve(x))

 def test_solve_dynamic(self):
   with self.test_session() as sess:
     for shape in self._shapes_to_test:
       for dtype in self._dtypes_to_test:
         operator, mat, feed_dict = self._operator_and_mat_and_feed_dict(
             shape, dtype, use_placeholder=True)
         rhs = self._make_rhs(operator)
         op_solve_v, mat_solve_v = sess.run(
             [operator.solve(rhs), tf.matrix_solve(mat, rhs)],
             feed_dict=feed_dict)
         self.assertAllClose(op_solve_v, mat_solve_v)

 def _verifySolve(self, x, y):
   for np_type in [np.float32, np.float64]:
     a = x.astype(np_type)
     b = y.astype(np_type)
     with self.test_session():
       if a.ndim == 2:
         tf_ans = tf.matrix_solve(a, b)
       else:
         tf_ans = tf.batch_matrix_solve(a, b)
       out = tf_ans.eval()
     np_ans = np.linalg.solve(a, b)
     self.assertEqual(np_ans.shape, out.shape)
     self.assertAllClose(np_ans, out)

def Bound2(phi_0,phi_1,phi_2,sigma_noise,K_mm,mean_y):
    # Preliminary Bound
    beta=1/tf.square(sigma_noise)
    bound=0
    N=h.get_dim(mean_y,0)
    M=h.get_dim(K_mm,0)
    W_inv_part=beta*phi_2+K_mm
    global phi_200
    phi_200=tf.matrix_solve(W_inv_part,tf.transpose(phi_1))
    W=beta*np.eye(N)-tf.square(beta)*h.Mul(phi_1,tf.matrix_solve(W_inv_part,tf.transpose(phi_1)))
    # Computations
    bound+=N*tf.log(beta)
    bound+=h.log_det(K_mm+1e-3*np.eye(M))
    bound-=h.Mul(tf.transpose(mean_y),W,mean_y)
    global matrix_determinant
    matrix_determinant=tf.ones(1) #h.log_det(W_inv_part+1e2*np.eye(M))#-1e-40*tf.exp(h.log_det(W_inv_part))


    bound-=h.log_det(W_inv_part+1e-3*tf.reduce_mean(W_inv_part)*np.eye(M))
    bound-=beta*phi_0
    bound+=beta*tf.trace(tf.cholesky_solve(tf.cholesky(K_mm),phi_2))
    bound=bound*0.5
    return bound

 def test_solve(self):
   with self.test_session() as sess:
     for shape in self._shapes_to_test:
       for dtype in self._dtypes_to_test:
         operator, mat, _ = self._operator_and_mat_and_feed_dict(
             shape, dtype, use_placeholder=False)
         for adjoint in [False, True]:
           if adjoint and operator.is_self_adjoint:
             continue
           rhs = self._make_rhs(operator)
           op_solve = operator.solve(rhs, adjoint=adjoint)
           mat_solve = tf.matrix_solve(self._maybe_adjoint(mat, adjoint), rhs)
           self.assertAllEqual(op_solve.get_shape(), mat_solve.get_shape())
           op_solve_v, mat_solve_v = sess.run([op_solve, mat_solve])
           self.assertAllClose(op_solve_v, mat_solve_v)

  def test_sqrt_solve(self):
    # Square roots are not unique, but we should still have
    # S^{-T} S^{-1} x = A^{-1} x.
    # In our case, we should have S = S^T, so then S^{-1} S^{-1} x = A^{-1} x.
    with self.test_session():
      for batch_shape in [(), (2, 3,)]:
        for k in [1, 4]:
          operator, mat = self._build_operator_and_mat(batch_shape, k)

          # Work with 5 simultaneous systems.  5 is arbitrary.
          x = self._rng.randn(*(batch_shape + (k, 5)))

          self._compare_results(
              expected=tf.matrix_solve(mat, x).eval(),
              actual=operator.sqrt_solve(operator.sqrt_solve(x)))

 def test_solve(self):
   self._maybe_skip("solve")
   with self.test_session() as sess:
     for use_placeholder in False, True:
       for shape in self._shapes_to_test:
         for dtype in self._dtypes_to_test:
           for adjoint in False, True:
             operator, mat, feed_dict = self._operator_and_mat_and_feed_dict(
                 shape, dtype, use_placeholder=use_placeholder)
             rhs = self._make_rhs(operator, adjoint=adjoint)
             op_solve = operator.solve(rhs, adjoint=adjoint)
             mat_solve = tf.matrix_solve(mat, rhs, adjoint=adjoint)
             if not use_placeholder:
               self.assertAllEqual(op_solve.get_shape(), mat_solve.get_shape())
             op_solve_v, mat_solve_v = sess.run(
                 [op_solve, mat_solve], feed_dict=feed_dict)
             self.assertAC(op_solve_v, mat_solve_v)