| 1052 | } |
| 1053 | |
| 1054 | PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim) { |
| 1055 | if (!PyList_Check(tensors) && !PyTuple_Check(tensors)) { |
| 1056 | PyErr_SetString(PyExc_TypeError, |
| 1057 | tensorflow::strings::StrCat( |
| 1058 | "tensors argument must be a list or a tuple. Got \"", |
| 1059 | Py_TYPE(tensors)->tp_name, "\"") |
| 1060 | .c_str()); |
| 1061 | return nullptr; |
| 1062 | } |
| 1063 | if (slice_dim < 0) { |
| 1064 | PyErr_SetString( |
| 1065 | PyExc_ValueError, |
| 1066 | tensorflow::strings::StrCat("Slice dimension must be non-negative. " |
| 1067 | "Got ", |
| 1068 | slice_dim) |
| 1069 | .c_str()); |
| 1070 | return nullptr; |
| 1071 | } |
| 1072 | |
| 1073 | Py_ssize_t num_tensors = PySequence_Fast_GET_SIZE(tensors); |
| 1074 | int64_t num_tensors_int = static_cast<int64_t>(num_tensors); |
| 1075 | auto tensor = tensorflow::make_safe(TF_AllocateTensor( |
| 1076 | TF_INT32, &num_tensors_int, /*num_dims=*/1, /*len=*/4 * num_tensors_int)); |
| 1077 | int32_t* data = reinterpret_cast<int32_t*>(TF_TensorData(tensor.get())); |
| 1078 | auto status = tensorflow::make_safe(TF_NewStatus()); |
| 1079 | for (Py_ssize_t i = 0; i < num_tensors; ++i) { |
| 1080 | PyObject* tensor_obj = PySequence_Fast_GET_ITEM(tensors, i); |
| 1081 | if (!EagerTensor_CheckExact(tensor_obj)) { |
| 1082 | PyErr_SetString(PyExc_TypeError, |
| 1083 | tensorflow::strings::StrCat( |
| 1084 | "Expected a list of EagerTensors but " |
| 1085 | "element ", |
| 1086 | i, " has type \"", Py_TYPE(tensor_obj)->tp_name, "\"") |
| 1087 | .c_str()); |
| 1088 | return nullptr; |
| 1089 | } |
| 1090 | |
| 1091 | EagerTensor* t = reinterpret_cast<EagerTensor*>(tensor_obj); |
| 1092 | TFE_TensorHandle* handle = t->handle; |
| 1093 | int num_dims = TFE_TensorHandleNumDims(handle, status.get()); |
| 1094 | if (MaybeRaiseExceptionFromTFStatus(status.get(), PyExc_ValueError)) { |
| 1095 | return nullptr; |
| 1096 | } |
| 1097 | if (slice_dim >= num_dims) { |
| 1098 | PyErr_SetString( |
| 1099 | PyExc_IndexError, |
| 1100 | tensorflow::strings::StrCat("Slice dimension (", slice_dim, |
| 1101 | ") must be smaller than rank of all " |
| 1102 | "tensors, but tensor at index ", |
| 1103 | i, " has rank ", num_dims) |
| 1104 | .c_str()); |
| 1105 | return nullptr; |
| 1106 | } |
| 1107 | int64_t dim = TFE_TensorHandleDim(handle, slice_dim, status.get()); |
| 1108 | if (MaybeRaiseExceptionFromTFStatus(status.get(), PyExc_ValueError)) { |
| 1109 | return nullptr; |
| 1110 | } |
| 1111 | data[i] = dim; |
nothing calls this directly
no test coverage detected