// CHECK: tf_mlrt.executeop{{.*}}op: \22MapDataset\22
  // CHECK-SAME: \22__inference_Dataset_map_lambda_16\22

  // MapDatasetOp's relationship with its associated function is as
  // follows: first M function params are dictated by the set
  // output shapes and types, the next N are the last Ninputs from MapDataset
  // op. The MapDataset op always has N+1 inputs.
  // TODO(jpienaar): Avoid this lookup.
  auto module = op->getParentOfType<ModuleOp>();
  auto f = module.lookupSymbol<func::FuncOp>(op.getF());

    // For this test case the first 12 inputs to the function f below will be
    // from ParallelMapDataset, while the remaining one is a captured input.
    %75 = "tf.MapDataset"(%74, %arg3) {
      device = "/job:localhost/replica:0/task:0/device:CPU:0",
      f = @__inference_Dataset_map_lambda_96150,
      f._tf_data_function = true,

}

def TF_FlatMapDatasetOp : TF_Op<"FlatMapDataset", [Pure]> {
  let summary = [{
Creates a dataset that applies `f` to the outputs of `input_dataset`.
  }];

  let description = [{
Unlike MapDataset, the `f` in FlatMapDataset is expected to return a
Dataset variant, and FlatMapDataset will flatten successive results
into a single Dataset.
  }];

  let arguments = (ins