Searched refs:main_diagonal (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | tridiagonal.cc | 51 XlaOp main_diagonal, in CheckSystemAndReturnNumEquations() argument
59 builder->GetShape(main_diagonal)); in CheckSystemAndReturnNumEquations()
145 StatusOr<XlaOp> ThomasSolver(XlaOp lower_diagonal, XlaOp main_diagonal, in ThomasSolver() argument
151 lower_diagonal, main_diagonal, upper_diagonal, rhs)); in ThomasSolver()
153 XlaOp main_diag_after_elimination = ZerosLike(main_diagonal); in ThomasSolver()
160 UpdateEq(main_diag_after_elimination, 0, Coefficient(main_diagonal, 0)); in ThomasSolver()
187 auto main_diagonal = values[1]; in ThomasSolver() local
196 auto main_diagonal_i = Coefficient(main_diagonal, i); in ThomasSolver()
214 main_diagonal, in ThomasSolver()
224 {lower_diagonal, main_diagonal, rhs, main_diag_after_elimination, in ThomasSolver()
[all …]
|
| D | tridiagonal_test.cc | 48 Array3D<float> main_diagonal(batch_size, 1, num_eqs); in XLA_TEST_P() local
53 main_diagonal.FillRandom(0.05, /*mean=*/1.0, in XLA_TEST_P()
67 main_diagonal, 1, "main_diagonal", &builder, &main_diagonal_xla); in XLA_TEST_P()
|
| D | tridiagonal.h | 25 StatusOr<XlaOp> ThomasSolver(XlaOp lower_diagonal, XlaOp main_diagonal,
|