| 149 | explicit UnwrapDatasetVariantOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} |
| 150 | |
| 151 | void Compute(OpKernelContext* ctx) override { |
| 152 | const Tensor& tensor = ctx->input(0); |
| 153 | OP_REQUIRES(ctx, |
| 154 | tensor.dtype() == DT_VARIANT && |
| 155 | TensorShapeUtils::IsScalar(tensor.shape()), |
| 156 | errors::InvalidArgument( |
| 157 | "Dataset tensor must be a scalar of dtype DT_VARIANT.")); |
| 158 | Variant variant = tensor.scalar<Variant>()(); |
| 159 | const WrappedDatasetVariantWrapper* wrapper = |
| 160 | variant.get<WrappedDatasetVariantWrapper>(); |
| 161 | OP_REQUIRES(ctx, wrapper != nullptr, |
| 162 | errors::InvalidArgument( |
| 163 | "Tensor must be a WrappedDataset variant object.")); |
| 164 | Tensor ds_tensor = wrapper->get(); |
| 165 | OP_REQUIRES_OK(ctx, ctx->set_output("output_handle", ds_tensor)); |
| 166 | } |
| 167 | }; |
| 168 | |
| 169 | REGISTER_KERNEL_BUILDER(Name("UnwrapDatasetVariant").Device(DEVICE_CPU), |
nothing calls this directly
no test coverage detected