| 208 | } |
| 209 | |
| 210 | void ParallelMapDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input, |
| 211 | DatasetBase** output) { |
| 212 | int32 num_parallel_calls; |
| 213 | OP_REQUIRES_OK( |
| 214 | ctx, ParseScalarArgument(ctx, kNumParallelCalls, &num_parallel_calls)); |
| 215 | OP_REQUIRES( |
| 216 | ctx, num_parallel_calls > 0 || num_parallel_calls == model::kAutotune, |
| 217 | errors::InvalidArgument("num_parallel_calls must be greater than zero.")); |
| 218 | |
| 219 | std::unique_ptr<CapturedFunction> captured_func; |
| 220 | OP_REQUIRES_OK(ctx, |
| 221 | CapturedFunction::Create(ctx, func_metadata_, kOtherArguments, |
| 222 | &captured_func)); |
| 223 | |
| 224 | if (num_parallel_calls == model::kAutotune) { |
| 225 | metrics::RecordTFDataAutotune(kDatasetType); |
| 226 | } |
| 227 | |
| 228 | *output = |
| 229 | new Dataset(ctx, input, num_parallel_calls, output_types_, output_shapes_, |
| 230 | sloppy_, std::move(captured_func), preserve_cardinality_); |
| 231 | } |
| 232 | |
| 233 | namespace { |
| 234 | REGISTER_KERNEL_BUILDER(Name("ParallelMapDataset").Device(DEVICE_CPU), |
nothing calls this directly
no test coverage detected