Home
last modified time | relevance | path

Searched refs:lower_diagonal (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dtridiagonal.cc50 StatusOr<int64> CheckSystemAndReturnNumEquations(XlaOp lower_diagonal, in CheckSystemAndReturnNumEquations() argument
54 XlaBuilder* builder = lower_diagonal.builder(); in CheckSystemAndReturnNumEquations()
57 builder->GetShape(lower_diagonal)); in CheckSystemAndReturnNumEquations()
145 StatusOr<XlaOp> ThomasSolver(XlaOp lower_diagonal, XlaOp main_diagonal, in ThomasSolver() argument
147 XlaBuilder* builder = lower_diagonal.builder(); in ThomasSolver()
151 lower_diagonal, main_diagonal, upper_diagonal, rhs)); in ThomasSolver()
186 auto lower_diagonal = values[0]; in ThomasSolver() local
195 auto lower_diagonal_i = Coefficient(lower_diagonal, i); in ThomasSolver()
213 return std::vector<XlaOp>{lower_diagonal, in ThomasSolver()
224 {lower_diagonal, main_diagonal, rhs, main_diag_after_elimination, in ThomasSolver()
[all …]
Dtridiagonal_test.cc47 Array3D<float> lower_diagonal(batch_size, 1, num_eqs); in XLA_TEST_P() local
52 lower_diagonal.FillRandom(1.0, /*mean=*/0.0, /*seed=*/0); in XLA_TEST_P()
65 lower_diagonal, 0, "lower_diagonal", &builder, &lower_diagonal_xla); in XLA_TEST_P()
Dtridiagonal.h25 StatusOr<XlaOp> ThomasSolver(XlaOp lower_diagonal, XlaOp main_diagonal,