| 282 | } |
| 283 | |
| 284 | Status ReadDiagIndex(InferenceContext* c, const Tensor* diag_index_tensor, |
| 285 | int32* lower_diag_index, int32* upper_diag_index) { |
| 286 | // This function assumes that the shape of diag_index_tensor is fully defined. |
| 287 | if (diag_index_tensor->dims() == 0) { |
| 288 | *lower_diag_index = diag_index_tensor->scalar<int32>()(); |
| 289 | *upper_diag_index = *lower_diag_index; |
| 290 | } else { |
| 291 | int32 num_elements = diag_index_tensor->dim_size(0); |
| 292 | if (num_elements == 1) { |
| 293 | *lower_diag_index = diag_index_tensor->vec<int32>()(0); |
| 294 | *upper_diag_index = *lower_diag_index; |
| 295 | } else if (num_elements == 2) { |
| 296 | *lower_diag_index = diag_index_tensor->vec<int32>()(0); |
| 297 | *upper_diag_index = diag_index_tensor->vec<int32>()(1); |
| 298 | } else { |
| 299 | return errors::InvalidArgument( |
| 300 | "diag_index must be a vector with one or two elements. It has ", |
| 301 | num_elements, " elements."); |
| 302 | } |
| 303 | } |
| 304 | return Status::OK(); |
| 305 | } |
| 306 | |
| 307 | } // namespace |
| 308 |
no test coverage detected