| 338 | } |
| 339 | |
| 340 | Status GetDatasetFromVariantTensor(const Tensor& tensor, |
| 341 | DatasetBase** out_dataset) { |
| 342 | if (!(tensor.dtype() == DT_VARIANT && |
| 343 | TensorShapeUtils::IsScalar(tensor.shape()))) { |
| 344 | return errors::InvalidArgument( |
| 345 | "Dataset tensor must be a scalar of dtype DT_VARIANT."); |
| 346 | } |
| 347 | const Variant& variant = tensor.scalar<Variant>()(); |
| 348 | const DatasetVariantWrapper* wrapper = variant.get<DatasetVariantWrapper>(); |
| 349 | if (wrapper == nullptr) { |
| 350 | return errors::InvalidArgument("Tensor must be a Dataset object."); |
| 351 | } |
| 352 | *out_dataset = wrapper->get(); |
| 353 | if (*out_dataset == nullptr) { |
| 354 | return errors::Internal("Read uninitialized Dataset variant."); |
| 355 | } |
| 356 | return Status::OK(); |
| 357 | } |
| 358 | |
| 359 | Status StoreDatasetInVariantTensor(DatasetBase* dataset, Tensor* tensor) { |
| 360 | if (!(tensor->dtype() == DT_VARIANT && |