Searched refs:diag_index_shape (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | matrix_diag_ops.cc | 59 TensorShape diag_index_shape = context->InputShape("k"); in ProcessDiagIndex() local
64 if (diag_index_shape.dims() == 0) { in ProcessDiagIndex()
|
| /external/tensorflow/tensorflow/core/framework/ |
| D | common_shape_fns.cc | 1360 ShapeHandle input_shape, diag_index_shape, unused_shape; in MatrixDiagPartV2Shape() local
1362 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &diag_index_shape)); in MatrixDiagPartV2Shape()
1366 if (!c->RankKnown(input_shape) || !c->FullyDefined(diag_index_shape) || in MatrixDiagPartV2Shape()
1414 ShapeHandle input_shape, diag_index_shape, unused_shape; in MatrixDiagV2Shape() local
1416 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &diag_index_shape)); in MatrixDiagV2Shape()
1423 if (!c->RankKnown(input_shape) || !c->FullyDefined(diag_index_shape) || in MatrixDiagV2Shape()
1514 ShapeHandle input_shape, diag_shape, diag_index_shape; in MatrixSetDiagV2Shape() local
1517 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(2), 1, &diag_index_shape)); in MatrixSetDiagV2Shape()
1523 if (diag_index_tensor != nullptr && c->FullyDefined(diag_index_shape)) { in MatrixSetDiagV2Shape()
|