Home
last modified time | relevance | path

Searched refs:max_diag_len (Results 1 – 18 of 18) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/linalg/
Dmatrix_diag_op_gpu.cu.cc30 const int max_diag_len, in ComputeContentOffset() argument
41 return max_diag_len - diag_len; in ComputeContentOffset()
47 const int num_diags, const int max_diag_len, const int lower_diag_index, in MatrixDiagKernel() argument
59 ComputeContentOffset(diag_index, max_diag_len, num_rows, num_cols, in MatrixDiagKernel()
64 diag_ptr[batch * num_diags * max_diag_len + in MatrixDiagKernel()
65 diag_index_in_input * max_diag_len + index_in_the_diagonal]; in MatrixDiagKernel()
79 const Eigen::Index max_diag_len, const T padding_value, in Compute()
86 if (batch_size == 0 || max_diag_len == 0 || num_rows == 0 || in Compute()
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
[all …]
Dmatrix_set_diag_op_gpu.cu.cc31 const int max_diag_len, in ComputeContentOffset() argument
42 return max_diag_len - diag_len; in ComputeContentOffset()
48 const int max_diag_len, const int upper_diag_index, in MatrixSetDiagKernel() argument
52 const int batch_and_diag_index = index / max_diag_len; in MatrixSetDiagKernel()
53 int index_in_the_diagonal = index - batch_and_diag_index * max_diag_len; in MatrixSetDiagKernel()
58 ComputeContentOffset(diag_index, max_diag_len, m, n, in MatrixSetDiagKernel()
74 const int max_diag_len, const int lower_diag_index, in MatrixCopyInputAndSetDiagKernel() argument
87 ComputeContentOffset(diag_index, max_diag_len, m, n, in MatrixCopyInputAndSetDiagKernel()
91 diag_ptr[batch * num_diags * max_diag_len + in MatrixCopyInputAndSetDiagKernel()
92 diag_index_in_input * max_diag_len + index_in_the_diagonal]; in MatrixCopyInputAndSetDiagKernel()
[all …]
Dmatrix_diag_op.cc128 const int32 max_diag_len = in Compute() local
131 output_shape.AddDim(max_diag_len); in Compute()
139 output_reshaped, lower_diag_index, upper_diag_index, max_diag_len, in Compute()
220 const Eigen::Index max_diag_len = diagonal_shape.dim_size(diag_rank - 1); in Compute() local
221 const int32 min_num_rows = max_diag_len - std::min(upper_diag_index, 0); in Compute()
222 const int32 min_num_cols = max_diag_len + std::max(lower_diag_index, 0); in Compute()
258 output_reshaped, lower_diag_index, upper_diag_index, max_diag_len, in Compute()
319 int diag_index, int max_diag_len, int num_rows, int num_cols, in ComputeDiagLenAndContentOffset() argument
325 const int content_offset = (left_align) ? 0 : (max_diag_len - diag_len); in ComputeDiagLenAndContentOffset()
336 const Eigen::Index max_diag_len, const T padding_value, in Compute()
[all …]
Dmatrix_set_diag_op.cc137 const int32 max_diag_len = in Compute() local
140 expected_diag_shape.AddDim(max_diag_len); in Compute()
165 output_reshaped, lower_diag_index, upper_diag_index, max_diag_len, in Compute()
210 const Eigen::Index max_diag_len, in Compute()
217 auto compute_shard = [&output, &diag, &upper_diag_index, &max_diag_len, in Compute()
223 Eigen::Index diag_base_index = begin * num_diags * max_diag_len; in Compute()
229 diag_index, max_diag_len, num_rows, num_cols, in Compute()
244 diag_base_index += max_diag_len; in Compute()
251 const Eigen::Index cost_per_batch = 10 * num_diags * max_diag_len; in Compute()
271 const Eigen::Index upper_diag_index, const Eigen::Index max_diag_len, \
Dmatrix_diag_op.h44 int diag_index, int max_diag_len, int num_rows, int num_cols,
53 const Eigen::Index upper_diag_index, const Eigen::Index max_diag_len,
65 const Eigen::Index max_diag_len, const T padding_value,
Dmatrix_set_diag_op.h34 const Eigen::Index max_diag_len,
/external/tensorflow/tensorflow/core/api_def/base_api/
Dapi_def_MatrixDiagPartV3.pbtxt30 Some diagonals are shorter than `max_diag_len` and need to be padded. `align` is
45 Let `max_diag_len` be the maximum length among all diagonals to be extracted,
46 `max_diag_len = min(M + min(k[1], 0), N + min(-k[0], 0))`
51 `[I, J, ..., L, max_diag_len]` and values:
61 `[I, J, ..., L, num_diags, max_diag_len]` with values:
72 offset = max_diag_len - diag_len(d) ; if (`align` in {RIGHT_LEFT, RIGHT_RIGHT}
123 # max_diag_len can be shorter than the main diagonal.
Dapi_def_MatrixSetDiagV3.pbtxt32 Some diagonals are shorter than `max_diag_len` and need to be padded. `align` is
48 `k[0] == k[1]`, `diagonal` has `r` dimensions `[I, J, ..., L, max_diag_len]`.
49 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
51 `max_diag_len` is the longest diagonal in the range `[k[0], k[1]]`,
52 `max_diag_len = min(M + min(k[1], 0), N + min(-k[0], 0))`
75 offset = max_diag_len - diag_len(d) ; if (`align` in {RIGHT_LEFT, RIGHT_RIGHT}
Dapi_def_MatrixSetDiagV2.pbtxt36 `k[0] == k[1]`, `diagonal` has `r` dimensions `[I, J, ..., L, max_diag_len]`.
37 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
39 `max_diag_len` is the longest diagonal in the range `[k[0], k[1]]`,
40 `max_diag_len = min(M + min(k[1], 0), N + min(-k[0], 0))`
Dapi_def_MatrixDiagPartV2.pbtxt33 Let `max_diag_len` be the maximum length among all diagonals to be extracted,
34 `max_diag_len = min(M + min(k[1], 0), N + min(-k[0], 0))`
39 `[I, J, ..., L, max_diag_len]` and values:
49 `[I, J, ..., L, num_diags, max_diag_len]` with values:
Dapi_def_MatrixDiagV3.pbtxt48 Some diagonals are shorter than `max_diag_len` and need to be padded. `align` is
96 offset = max_diag_len - diag_len(d) ; if (`align` in {RIGHT_LEFT, RIGHT_RIGHT}
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dmatrix_diag_ops.cc120 const int64 upper_diag_index, const int64 max_diag_len, in SetMatrixDiag() argument
204 int padding_low = length_to_pad_to - max_diag_len; in SetMatrixDiag()
274 const int64 max_diag_len = diag_shape.dim_size(diag_rank - 1); in Compile() local
283 max_diag_len - std::min(upper_diag_index, int64{0}); in Compile()
285 max_diag_len + std::max(lower_diag_index, int64{0}); in Compile()
319 lower_diag_index, upper_diag_index, max_diag_len, in Compile()
392 const int32 max_diag_len = in Compile() local
395 output_shape.AddDim(max_diag_len); in Compile()
414 const int64 padding_len = max_diag_len - diag_len; in Compile()
507 const int32 max_diag_len = in Compile() local
[all …]
/external/tensorflow/tensorflow/compiler/tests/
Dmatrix_diag_ops_test.py52 max_diag_len = packed_diagonals.shape[-1]
60 padding_len = max_diag_len - diag_len
68 padding_len, max_diag_len, 1),)
/external/tensorflow/tensorflow/python/ops/
Darray_grad.py468 max_diag_len = math_ops.minimum(matrix_shape[0] + y_offset,
474 lambda: ops.convert_to_tensor([max_diag_len]),
476 max_diag_len]))
505 max_diag_len = math_ops.minimum(matrix_shape[0] + y_offset,
511 lambda: ops.convert_to_tensor([max_diag_len]),
513 max_diag_len]))
/external/tensorflow/tensorflow/python/kernel_tests/
Ddiag_op_test.py57 max_diag_len = packed_diagonals.shape[-1]
65 padding_len = max_diag_len - diag_len
73 padding_len, max_diag_len, 1),)
/external/tensorflow/tensorflow/core/framework/
Dcommon_shape_fns.cc1293 int32 max_diag_len = InferenceContext::kUnknownDim; in MatrixDiagPartV2Shape() local
1304 max_diag_len = std::min(num_rows + std::min(upper_diag_index, 0), in MatrixDiagPartV2Shape()
1316 dims.push_back(c->MakeDim(max_diag_len)); in MatrixDiagPartV2Shape()
1377 const int32 max_diag_len = c->Value(c->Dim(input_shape, input_rank - 1)); in MatrixDiagV2Shape() local
1378 const int32 min_num_rows = max_diag_len - std::min(upper_diag_index, 0); in MatrixDiagV2Shape()
1379 const int32 min_num_cols = max_diag_len + std::max(lower_diag_index, 0); in MatrixDiagV2Shape()
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf.cc1520 int64_t max_diag_len = in matchAndRewrite() local
1526 Shape indices_shape({1, num_diags, max_diag_len}); in matchAndRewrite()
1543 BroadcastConstant(loc, indices_shape, max_diag_len, 32, rewriter); 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()
1627 output_shape.push_back(max_diag_len); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_generated_ops.td7248 Let `max_diag_len` be the maximum length among all diagonals to be extracted,
7249 `max_diag_len = min(M + min(k[1], 0), N + min(-k[0], 0))`
7254 `[I, J, ..., L, max_diag_len]` and values:
7264 `[I, J, ..., L, num_diags, max_diag_len]` with values:
7275 offset = max_diag_len - diag_len(d) ; if (`align` in {RIGHT_LEFT, RIGHT_RIGHT}
7326 # max_diag_len can be shorter than the main diagonal.
7524 offset = max_diag_len - diag_len(d) ; if (`align` in {RIGHT_LEFT, RIGHT_RIGHT}
7710 `k[0] == k[1]`, `diagonal` has `r` dimensions `[I, J, ..., L, max_diag_len]`.
7711 Otherwise, it has `r+1` dimensions `[I, J, ..., L, num_diags, max_diag_len]`.
7713 `max_diag_len` is the longest diagonal in the range `[k[0], k[1]]`,
[all …]