| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | matrix_solve_ls_op_test.py | 46 def _GenerateTestData(matrix_shape, num_rhs): argument
54 rhs = np.ones([m, num_rhs]).astype(np.float32)
237 num_rhs = 1
242 rhs = np.ones([matrix_shape[0], num_rhs]).astype(np.float32)
254 num_rhs = 1
259 rhs = np.ones([matrix_shape[0], num_rhs]).astype(np.float32)
271 num_rhs = 1
276 rhs = np.ones([matrix_shape[0], num_rhs]).astype(np.float32)
331 for num_rhs in 1, 2, matrix_shape[-1]:
336 matrix, rhs = _GenerateTestData(matrix_shape, num_rhs)
[all …]
|
| D | matrix_solve_op_test.py | 161 def _GenerateTestData(self, matrix_shape, num_rhs): argument
168 rhs = np.ones([n, num_rhs]).astype(np.float32)
179 for num_rhs in 1, 2, matrix_shape[-1]:
184 matrix, rhs = self._GenerateTestData(matrix_shape, num_rhs)
195 num_rhs=num_rhs,
202 matrix, rhs = self._GenerateTestData(matrix_shape, num_rhs)
212 matrix_shape=matrix_shape, num_rhs=num_rhs,
|
| D | tridiagonal_solve_op_test.py | 670 def _generateData(self, matrix_size, batch_size, num_rhs, seed=42): argument
672 data = np.random.normal(size=(batch_size, matrix_size, 3 + num_rhs))
678 def _generateMatrixData(self, matrix_size, batch_size, num_rhs, seed=42): argument
686 vector = np.random.randn(batch_size, matrix_size, num_rhs)
700 matrix_size, batch_size, num_rhs = size_option
705 diags, rhs = generate_data_fn(matrix_size, batch_size, num_rhs)
718 batch_size, num_rhs,
|
| /external/tensorflow/tensorflow/core/kernels/linalg/ |
| D | tridiagonal_solve_op_gpu.cu.cc | 46 const int num_rhs, in SolveForSizeOneOrTwoKernel() argument
54 for (int i : GpuGridRangeX(num_rhs)) { in SolveForSizeOneOrTwoKernel()
63 for (int i : GpuGridRangeX(num_rhs)) { in SolveForSizeOneOrTwoKernel()
64 x[i] = (diags[3] * rhs[i] - diags[0] * rhs[i + num_rhs]) / det; in SolveForSizeOneOrTwoKernel()
65 x[i + num_rhs] = (diags[2] * rhs[i + num_rhs] - diags[5] * rhs[i]) / det; in SolveForSizeOneOrTwoKernel()
202 const int num_rhs) const { in SolveWithGtsv()
208 num_eqs, num_rhs, subdiag, diag, superdiag, rhs, in SolveWithGtsv()
219 num_eqs, num_rhs, subdiag, diag, superdiag, rhs, in SolveWithGtsv()
257 const int64 num_rhs = rhs.dim_size(rhs.dims() - 1); in Compute() local
268 || num_rhs > 1 // GtsvBatched doesn't support multiple rhs in Compute()
|
| /external/tensorflow/tensorflow/compiler/tests/ |
| D | tridiagonal_solve_ops_test.py | 98 num_rhs = 5
103 num_rhs)).astype(np.float32)
109 shape=(batch_size, num_dims, num_rhs), dtype=dtypes.float32)
121 for eq in range(num_rhs):
220 num_rhs = 5
225 num_rhs)).astype(np.float32)
234 shape=(batch_size, num_dims, num_rhs), dtype=dtypes.float32)
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | eigen_contraction_kernel.h | 599 allocateSlices(Device& d, const int num_lhs, const int num_rhs, \
607 d, bm, bk, bn, num_lhs, num_rhs, num_slices, lhs_mem.data(), \
615 if (num_rhs > 0) rhs_blocks[x].resize(num_rhs); \
616 for (Index n = 0; n < num_rhs; n++) { \
794 allocateSlices(Device& d, const int num_lhs, const int num_rhs, \
802 d, bm, bk, bn, num_lhs, num_rhs, num_slices, lhs_mem.data(), \
810 if (num_rhs > 0) rhs_blocks[x].resize(num_rhs); \
811 for (Index n = 0; n < num_rhs; n++) { \
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | tridiagonal_test.cc | 45 const int64 num_rhs = std::get<2>(spec); in XLA_TEST_P() local
50 Array3D<float> rhs(batch_size, num_rhs, num_eqs); in XLA_TEST_P()
|