Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/linalg/
H A Dmatrix_diag_op.cc205 auto& num_rows_tensor = context->input(2); in Compute() local
206 OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_rows_tensor.shape()), in Compute()
208 num_rows = num_rows_tensor.flat<int32>()(0); in Compute()
/aosp_15_r20/external/tensorflow/tensorflow/core/framework/
H A Dcommon_shape_fns.cc1454 const Tensor* num_rows_tensor = c->input_tensor(2); in MatrixDiagV2Shape() local
1458 if (num_rows_tensor != nullptr) { in MatrixDiagV2Shape()
1459 TF_RETURN_IF_ERROR(c->GetScalarFromTensor(num_rows_tensor, &num_rows)); in MatrixDiagV2Shape()