Home
last modified time | relevance | path

Searched refs:num_cols_tensor (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/linalg/
H A Dmatrix_diag_op.cc210 auto& num_cols_tensor = context->input(3); in Compute() local
211 OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_cols_tensor.shape()), in Compute()
213 num_cols = num_cols_tensor.flat<int32>()(0); in Compute()
/aosp_15_r20/external/tensorflow/tensorflow/core/framework/
H A Dcommon_shape_fns.cc1455 const Tensor* num_cols_tensor = c->input_tensor(3); in MatrixDiagV2Shape() local
1461 if (num_cols_tensor != nullptr) { in MatrixDiagV2Shape()
1462 TF_RETURN_IF_ERROR(c->GetScalarFromTensor(num_cols_tensor, &num_cols)); in MatrixDiagV2Shape()