| 141 | } |
| 142 | |
| 143 | Status TensorShapeFromTensor(const Tensor& t, PartialTensorShape* out) { |
| 144 | if (t.shape() == TensorShape({})) { |
| 145 | if ((t.dtype() == DT_INT32 && t.scalar<int32>()() == -1) || |
| 146 | (t.dtype() == DT_INT64 && t.scalar<int64>()() == -1)) { |
| 147 | *out = PartialTensorShape(); |
| 148 | return Status::OK(); |
| 149 | } |
| 150 | return errors::InvalidArgument( |
| 151 | "The only valid scalar shape tensor is the fully unknown shape " |
| 152 | "specified as -1."); |
| 153 | } |
| 154 | if (t.dtype() == DT_INT32) { |
| 155 | return PartialTensorShape::MakePartialShape(t.vec<int32>().data(), |
| 156 | t.NumElements(), out); |
| 157 | } else if (t.dtype() == DT_INT64) { |
| 158 | return PartialTensorShape::MakePartialShape(t.vec<int64>().data(), |
| 159 | t.NumElements(), out); |
| 160 | } |
| 161 | return errors::InvalidArgument( |
| 162 | "Expected an int32 or int64 shape tensor; found ", |
| 163 | DataTypeString(t.dtype())); |
| 164 | } |
| 165 | |
| 166 | Status GetElementShapeFromInput(OpKernelContext* c, |
| 167 | const TensorList& tensor_list, int index, |
no test coverage detected