content: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import numpy as np
  import pytest  # noqa: F401
  import torch  # needed for import of Torch C++ extensions to work
  from scipy.sparse import csr_matrix
  
  from theseus.utils import random_sparse_matrix
  
  
  # ideally we would like to support batch_size <= init_batch_size, but
  # because of limitations of cublas those have to be always identical
  def check_lu_solver(
      init_batch_size, batch_size, num_rows, num_cols, fill, verbose=False
  ):
      # this is necessary assumption, so that the hessian is full rank
      assert num_rows >= num_cols
  
      if not torch.cuda.is_available():
          return
      from theseus.extlib.cusolver_lu_solver import CusolverLUSolver
  
      rng = torch.Generator()
      rng.manual_seed(0)
      A_col_ind, A_row_ptr, A_val, _ = random_sparse_matrix(
          batch_size, num_rows, num_cols, fill, 3, rng, "cuda:0"
      )
      A_num_cols = num_cols
      A_num_rows = A_row_ptr.size(0) - 1
  
      b = torch.rand((batch_size, A_num_rows), dtype=torch.double).cuda()
  
      A_csr = [
          csr_matrix(
              (A_val[i].cpu(), A_col_ind.cpu(), A_row_ptr.cpu()), (A_num_rows, A_num_cols)
          )
          for i in range(batch_size)
      ]
      if verbose:
          print("A[0]:\n", A_csr[0].todense())
          print("b[0]:\n", b[0])
  
      AtA_csr = [(a.T @ a).tocsr() for a in A_csr]
      AtA_row_ptr = torch.tensor(AtA_csr[0].indptr).cuda()
      AtA_col_ind = torch.tensor(AtA_csr[0].indices).cuda()
      AtA_val = torch.tensor(np.array([m.data for m in AtA_csr])).cuda()
      AtA_num_rows = AtA_row_ptr.size(0) - 1
      AtA_num_cols = AtA_num_rows
  
      if verbose:
          print("AtA[0]:\n", AtA_csr[0].todense())
  
      slv = CusolverLUSolver(init_batch_size, AtA_num_cols, AtA_row_ptr, AtA_col_ind)
      singularities = slv.factor(AtA_val)
  
      if verbose:
          print("singularities:", singularities)
  
      b = torch.rand((batch_size, A_num_rows), dtype=torch.double).cuda()
      Atb = torch.tensor(
          np.array([A_csr[i].T @ b[i].cpu().numpy() for i in range(batch_size)])
      ).cuda()
      if verbose:
          print("Atb[0]:", Atb[0])
  
      sol = Atb.clone()
      slv.solve(sol)
      if verbose:
          print("x[0]:", sol[0])
  
      residuals = [
          AtA_csr[i] @ sol[i].cpu().numpy() - Atb[i].cpu().numpy()
          for i in range(batch_size)
      ]
      if verbose:
          print("residual[0]:", residuals[0])
  
      assert all(np.linalg.norm(res) < 1e-10 for res in residuals)
  
  
  @pytest.mark.cudaext
  def test_lu_solver_1():
      check_lu_solver(init_batch_size=5, batch_size=5, num_rows=50, num_cols=30, fill=0.2)
  
  
  @pytest.mark.cudaext
  def test_lu_solver_2():
      check_lu_solver(
          init_batch_size=5, batch_size=5, num_rows=150, num_cols=60, fill=0.2
      )
  
  
  @pytest.mark.cudaext
  def test_lu_solver_3():
      check_lu_solver(
          init_batch_size=10, batch_size=10, num_rows=300, num_cols=90, fill=0.2
      )
  
  
  @pytest.mark.cudaext
  def test_lu_solver_4():
      check_lu_solver(init_batch_size=5, batch_size=5, num_rows=50, num_cols=30, fill=0.1)
  
  
  @pytest.mark.cudaext
  def test_lu_solver_5():
      check_lu_solver(
          init_batch_size=5, batch_size=5, num_rows=150, num_cols=60, fill=0.1
      )
  
  
  @pytest.mark.cudaext
  def test_lu_solver_6():
      check_lu_solver(
          init_batch_size=10, batch_size=10, num_rows=300, num_cols=90, fill=0.1
      )
  
  
  # would like to test when irregular batch_size < init_batch_size,
  # but this is currently not supported by cublas, maybe in the future
  # def test_lu_solver_7():
  #     check_lu_solver(init_batch_size=10, batch_size=5, num_rows=150, num_cols=60, fill=0.2)
  