| 961 | } |
| 962 | |
| 963 | Status InferenceContext::GetScalarFromTensor(const Tensor* t, int64* val) { |
| 964 | // Caller must ensure that <t> is not NULL. |
| 965 | const int rank = t->dims(); |
| 966 | if (rank != 0) { |
| 967 | return errors::InvalidArgument("Input must be scalar but has rank ", rank); |
| 968 | } |
| 969 | |
| 970 | if (t->dtype() == DT_INT32) { |
| 971 | *val = t->scalar<int32>()(); |
| 972 | return Status::OK(); |
| 973 | } else if (t->dtype() == DT_INT64) { |
| 974 | *val = t->scalar<int64>()(); |
| 975 | return Status::OK(); |
| 976 | } else { |
| 977 | return errors::InvalidArgument("Scalar input must be int32 or int64."); |
| 978 | } |
| 979 | } |
| 980 | |
| 981 | // Returns a new dimension whose value is given by a scalar input tensor. |
| 982 | Status InferenceContext::MakeDimForScalarInput(int idx, DimensionHandle* out) { |
no test coverage detected