Function ThomasSolver

tensorflow/compiler/xla/client/lib/tridiagonal.cc:135–179

Applies Thomas algorithm to solve a linear system where the linear operand is a tri-diagonal matrix. See https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm for a simple reference on the Thomas algorithm. It is expected that the three diagonals are represented as tensors of shape [..., 1, num_equations] where num_equations is the number of dimensions of the unknowns considered in the linear

// the right-hand-side `rhs` should be [..., num_rhs, num_equations]. The
// solution will have the shape [..., num_rhs, num_equations].
StatusOr<XlaOp> ThomasSolver(XlaOp lower_diagonal, XlaOp main_diagonal,
                             XlaOp upper_diagonal, XlaOp rhs) {
  TF_ASSIGN_OR_RETURN(TridiagonalSystemShape system_shape,
                      CheckSystemAndReturnShape(lower_diagonal, main_diagonal,
                                                upper_diagonal, rhs));

  auto rank = system_shape.rank;
  auto num_eqs = system_shape.num_equations;

  std::vector<XlaOp> main_diag_after_elimination(num_eqs);
  std::vector<XlaOp> rhs_after_elimination(num_eqs);
  std::vector<XlaOp> upper_diagonal_coeffs(num_eqs);

  main_diag_after_elimination[0] = Coefficient(main_diagonal, 0);
  rhs_after_elimination[0] = Coefficient(rhs, 0);
  for (int64 i = 0; i < num_eqs - 1; i++) {
    upper_diagonal_coeffs[i] = Coefficient(upper_diagonal, i);
  }

  // Forward transformation.
  for (int64 i = 1; i < num_eqs; i++) {
    auto lower_diagonal_i = Coefficient(lower_diagonal, i);
    auto main_diagonal_i = Coefficient(main_diagonal, i);
    auto rhs_i = Coefficient(rhs, i);

    auto w_i = lower_diagonal_i / main_diag_after_elimination[i - 1];

    main_diag_after_elimination[i] =
        main_diagonal_i - w_i * upper_diagonal_coeffs[i - 1];
    rhs_after_elimination[i] = rhs_i - w_i * rhs_after_elimination[i - 1];
  }

  std::vector<XlaOp> x_coeffs(num_eqs);

  // Backward reduction.
  x_coeffs[num_eqs - 1] = rhs_after_elimination[num_eqs - 1] /
                          main_diag_after_elimination[num_eqs - 1];
  for (int i = num_eqs - 2; i >= 0; i--) {
    x_coeffs[i] = (rhs_after_elimination[i] -
                   upper_diagonal_coeffs[i] * x_coeffs[i + 1]) /
                  main_diag_after_elimination[i];
  }

  return ConcatInDim(lower_diagonal.builder(), x_coeffs, rank - 1);
}

// Applies Thomas algorithm to solve a linear system where the linear operand
// is a tri-diagonal matrix.