| /external/tensorflow/tensorflow/core/kernels/linalg/ |
| D | matrix_set_diag_op_gpu.cu.cc | 47 const int num_threads, const int m, const int n, const int num_diags, in MatrixSetDiagKernel() argument
54 const int batch = batch_and_diag_index / num_diags; in MatrixSetDiagKernel()
55 const int diag_index_in_input = batch_and_diag_index - batch * num_diags; in MatrixSetDiagKernel()
73 const int num_threads, const int m, const int n, const int num_diags, in MatrixCopyInputAndSetDiagKernel() argument
91 diag_ptr[batch * num_diags * max_diag_len + in MatrixCopyInputAndSetDiagKernel()
113 const int num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
118 GetGpuLaunchConfig(batch_size * num_diags * max_diag_len, device); in Compute()
122 config.virtual_thread_count, m, n, num_diags, max_diag_len, in Compute()
130 config.virtual_thread_count, m, n, num_diags, max_diag_len, in Compute()
|
| D | matrix_diag_op_gpu.cu.cc | 47 const int num_diags, const int max_diag_len, const int lower_diag_index, in MatrixDiagKernel() argument
64 diag_ptr[batch * num_diags * max_diag_len + in MatrixDiagKernel()
85 const int num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
95 num_diags, max_diag_len, lower_diag_index, upper_diag_index, in Compute()
104 const int num_diags, const int max_diag_len, const int lower_diag_index, in MatrixDiagPartKernel() argument
112 const int batch = batch_and_mapped_diag_index / num_diags; in MatrixDiagPartKernel()
114 batch_and_mapped_diag_index - batch * num_diags; in MatrixDiagPartKernel()
146 const int num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
152 GetGpuLaunchConfig(batch_size * num_diags * max_diag_len, device); in Compute()
156 num_diags, max_diag_len, lower_diag_index, upper_diag_index, in Compute()
|
| D | matrix_set_diag_op.cc | 125 const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
129 (diag_shape.dim_size(input_rank - 2) == num_diags), in Compute()
136 if (num_diags > 1) expected_diag_shape.AddDim(num_diags); in Compute()
216 const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
218 &num_diags, &left_align_superdiagonal, in Compute()
223 Eigen::Index diag_base_index = begin * num_diags * max_diag_len; in Compute()
225 for (Eigen::Index m = 0; m < num_diags; ++m) { in Compute()
251 const Eigen::Index cost_per_batch = 10 * num_diags * max_diag_len; in Compute()
|
| D | matrix_diag_op.cc | 126 const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
127 if (num_diags > 1) output_shape.AddDim(num_diags); in Compute()
203 const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
215 diagonal_shape.dim_size(diag_rank - 2) == num_diags, in Compute()
244 if (num_diags == 1) { // Output has rank `rank+1`. in Compute()
351 const int num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
352 const int diag_elements_in_batch = num_diags * max_diag_len; in Compute()
397 const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1; in Compute() local
398 const Eigen::Index output_elements_in_batch = num_diags * max_diag_len; in Compute()
405 &upper_diag_index, &max_diag_len, &num_diags, in Compute()
[all …]
|
| D | tridiagonal_solve_op.cc | 56 auto num_diags = input_matrix_shapes[0].dim_size(0); in ValidateInputMatrixShapes() local
58 context, num_diags == 3, in ValidateInputMatrixShapes()
61 num_diags, " rows.")); in ValidateInputMatrixShapes()
|
| D | tridiagonal_solve_op_gpu.cu.cc | 113 auto num_diags = input_matrix_shapes[0].dim_size(0); in ValidateInputMatrixShapes() local
115 context, num_diags == 3, in ValidateInputMatrixShapes()
118 num_diags, " columns.")); in ValidateInputMatrixShapes()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | matrix_diag_ops.cc | 119 const int64 num_diags, const int64 lower_diag_index, in SetMatrixDiag() argument
165 if (num_diags > 1) { in SetMatrixDiag()
275 const int64 num_diags = upper_diag_index - lower_diag_index + 1; in Compile() local
278 num_diags == 1 || num_diags == diag_shape.dim_size(diag_rank - 2), in Compile()
312 output_shape.RemoveLastDims((num_diags == 1) ? 1 : 2); in Compile()
318 0, SetMatrixDiag(output, diag, output_shape, diag_rank, num_diags, in Compile()
390 const int num_diags = upper_diag_index - lower_diag_index + 1; in Compile() local
391 if (num_diags > 1) output_shape.AddDim(num_diags); in Compile()
402 if (num_diags == 1) { in Compile()
495 const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1; in Compile() local
[all …]
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_MatrixDiagPartV2.pbtxt | 35 Let `num_diags` be the number of diagonals to extract,
36 `num_diags = k[1] - k[0] + 1`.
38 If `num_diags == 1`, the output tensor is of rank `r - 1` with shape
49 `[I, J, ..., L, num_diags, max_diag_len]` with values:
|
| D | api_def_MatrixDiagPartV3.pbtxt | 47 Let `num_diags` be the number of diagonals to extract,
48 `num_diags = k[1] - k[0] + 1`.
50 If `num_diags == 1`, the output tensor is of rank `r - 1` with shape
61 `[I, J, ..., L, num_diags, max_diag_len]` with values:
|
| D | api_def_MatrixSetDiagV2.pbtxt | 37 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
38 `num_diags` is the number of diagonals, `num_diags = k[1] - k[0] + 1`.
|
| D | api_def_MatrixSetDiagV3.pbtxt | 49 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
50 `num_diags` is the number of diagonals, `num_diags = k[1] - k[0] + 1`.
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | diag_op_test.py | 747 num_diags = upper_diag_index - lower_diag_index + 1
748 num_diags_dim = () if num_diags == 1 else (num_diags,)
|
| /external/llvm-project/lldb/source/Plugins/ExpressionParser/Clang/ |
| D | ClangExpressionParser.cpp | 1262 size_t num_diags = diagnostics.size(); in RewriteExpression() local
1263 if (num_diags == 0) in RewriteExpression()
|
| /external/tensorflow/tensorflow/core/framework/ |
| D | common_shape_fns.cc | 1350 const int32 num_diags = c->Value(c->Dim(input_shape, input_rank - 2)); in MatrixDiagV2Shape() local
1353 if (num_diags != (upper_diag_index - lower_diag_index + 1)) { in MatrixDiagV2Shape()
1357 "num_diags = ", num_diags, ", d_lower = ", lower_diag_index, in MatrixDiagV2Shape()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 1515 int num_diags = k[1] - k[0] + 1; in matchAndRewrite() local
1526 Shape indices_shape({1, num_diags, max_diag_len}); in matchAndRewrite()
1609 RankedTensorType::get(Shape({num_diags, max_diag_len}), in matchAndRewrite()
1614 Shape concat_shape({2, num_diags, max_diag_len}); in matchAndRewrite()
1626 output_shape.push_back(num_diags); in matchAndRewrite()
1676 if (num_diags == 1) { in matchAndRewrite()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
| D | tf_generated_ops.td | 7250 Let `num_diags` be the number of diagonals to extract,
7251 `num_diags = k[1] - k[0] + 1`.
7253 If `num_diags == 1`, the output tensor is of rank `r - 1` with shape
7264 `[I, J, ..., L, num_diags, max_diag_len]` with values:
7711 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
7712 `num_diags` is the number of diagonals, `num_diags = k[1] - k[0] + 1`.
7809 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
7810 `num_diags` is the number of diagonals, `num_diags = k[1] - k[0] + 1`.
|