| 409 | } |
| 410 | |
| 411 | void DatasetOpKernel::Compute(OpKernelContext* ctx) { |
| 412 | DatasetBase* dataset = nullptr; |
| 413 | MakeDataset(ctx, &dataset); |
| 414 | if (ctx->status().ok()) { |
| 415 | Tensor* output = nullptr; |
| 416 | OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &output)); |
| 417 | OP_REQUIRES_OK(ctx, StoreDatasetInVariantTensor(dataset, output)); |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | void UnaryDatasetOpKernel::MakeDataset(OpKernelContext* ctx, |
| 422 | DatasetBase** output) { |