Searched refs:diag_index_tensor (Results 1 – 1 of 1) sorted by relevance
1245 Status ReadDiagIndex(InferenceContext* c, const Tensor* diag_index_tensor, in ReadDiagIndex() argument1248 if (diag_index_tensor->dims() == 0) { in ReadDiagIndex()1249 *lower_diag_index = diag_index_tensor->scalar<int32>()(); in ReadDiagIndex()1252 int32 num_elements = diag_index_tensor->dim_size(0); in ReadDiagIndex()1254 *lower_diag_index = diag_index_tensor->vec<int32>()(0); in ReadDiagIndex()1257 *lower_diag_index = diag_index_tensor->vec<int32>()(0); in ReadDiagIndex()1258 *upper_diag_index = diag_index_tensor->vec<int32>()(1); in ReadDiagIndex()1274 const Tensor* diag_index_tensor = c->input_tensor(1); in MatrixDiagPartV2Shape() local1276 diag_index_tensor == nullptr) { in MatrixDiagPartV2Shape()1282 TF_RETURN_IF_ERROR(ReadDiagIndex(c, diag_index_tensor, &lower_diag_index, in MatrixDiagPartV2Shape()[all …]