1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #include "tensorflow/core/framework/common_shape_fns.h" 16 #include "tensorflow/core/framework/op.h" 17 18 namespace tensorflow { 19 20 REGISTER_OP("AssertCardinalityDataset") 21 .Input("input_dataset: variant") 22 .Input("cardinality: int64") 23 .Output("handle: variant") 24 .Attr("output_types: list(type) >= 1") 25 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0102(shape_inference::InferenceContext* c) 26 .SetShapeFn([](shape_inference::InferenceContext* c) { 27 shape_inference::ShapeHandle unused; 28 // cardinality should be a scalar. 29 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 30 return shape_inference::ScalarShape(c); 31 }); 32 33 REGISTER_OP("AssertNextDataset") 34 .Input("input_dataset: variant") 35 .Input("transformations: string") 36 .Output("handle: variant") 37 .Attr("output_types: list(type) >= 1") 38 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0202(shape_inference::InferenceContext* c) 39 .SetShapeFn([](shape_inference::InferenceContext* c) { 40 shape_inference::ShapeHandle unused; 41 // transformations should be a vector. 42 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &unused)); 43 return shape_inference::ScalarShape(c); 44 }); 45 46 REGISTER_OP("ExperimentalAssertNextDataset") 47 .Input("input_dataset: variant") 48 .Input("transformations: string") 49 .Output("handle: variant") 50 .Attr("output_types: list(type) >= 1") 51 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0302(shape_inference::InferenceContext* c) 52 .SetShapeFn([](shape_inference::InferenceContext* c) { 53 shape_inference::ShapeHandle unused; 54 // transformations should be a vector. 55 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &unused)); 56 return shape_inference::ScalarShape(c); 57 }); 58 59 REGISTER_OP("AutoShardDataset") 60 .Input("input_dataset: variant") 61 .Input("num_workers: int64") 62 .Input("index: int64") 63 .Output("handle: variant") 64 .Attr("auto_shard_policy: int = 0") 65 .Attr("output_types: list(type) >= 1") 66 .Attr("output_shapes: list(shape) >= 1") 67 .Attr("num_replicas: int = 0") 68 .SetShapeFn(shape_inference::ScalarShape); 69 70 REGISTER_OP("ExperimentalAutoShardDataset") 71 .Input("input_dataset: variant") 72 .Input("num_workers: int64") 73 .Input("index: int64") 74 .Output("handle: variant") 75 .Attr("auto_shard_policy: int = 0") 76 .Attr("output_types: list(type) >= 1") 77 .Attr("output_shapes: list(shape) >= 1") 78 .SetShapeFn(shape_inference::ScalarShape); 79 80 REGISTER_OP("BytesProducedStatsDataset") 81 .Input("input_dataset: variant") 82 .Input("tag: string") 83 .Output("handle: variant") 84 .Attr("output_types: list(type) >= 1") 85 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0402(shape_inference::InferenceContext* c) 86 .SetShapeFn([](shape_inference::InferenceContext* c) { 87 shape_inference::ShapeHandle tag_shape; 88 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &tag_shape)); 89 return shape_inference::ScalarShape(c); 90 }); 91 92 REGISTER_OP("ExperimentalBytesProducedStatsDataset") 93 .Input("input_dataset: variant") 94 .Input("tag: string") 95 .Output("handle: variant") 96 .Attr("output_types: list(type) >= 1") 97 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0502(shape_inference::InferenceContext* c) 98 .SetShapeFn([](shape_inference::InferenceContext* c) { 99 shape_inference::ShapeHandle tag_shape; 100 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &tag_shape)); 101 return shape_inference::ScalarShape(c); 102 }); 103 104 REGISTER_OP("ChooseFastestBranchDataset") 105 .Input("input_dataset: variant") 106 .Input("ratio_numerator: int64") 107 .Input("ratio_denominator: int64") 108 .Input("other_arguments: Targuments") 109 .Output("handle: variant") 110 .Attr("Targuments: list(type) >= 0") 111 .Attr("num_elements_per_branch: int >= 1") 112 .Attr("branches: list(func) >= 1") 113 .Attr("other_arguments_lengths: list(int) >= 1") 114 .Attr("output_types: list(type) >= 1") 115 .Attr("output_shapes: list(shape) >= 1") 116 .SetShapeFn(shape_inference::ScalarShape); 117 118 REGISTER_OP("ChooseFastestDataset") 119 .Input("input_datasets: N * variant") 120 .Output("handle: variant") 121 .Attr("N: int >= 2") 122 .Attr("num_experiments: int") 123 .Attr("output_types: list(type) >= 1") 124 .Attr("output_shapes: list(shape) >= 1") 125 .SetShapeFn(shape_inference::ScalarShape); 126 127 REGISTER_OP("ExperimentalChooseFastestDataset") 128 .Input("input_datasets: N * variant") 129 .Output("handle: variant") 130 .Attr("N: int >= 2") 131 .Attr("num_experiments: int") 132 .Attr("output_types: list(type) >= 1") 133 .Attr("output_shapes: list(shape) >= 1") 134 .SetShapeFn(shape_inference::ScalarShape); 135 136 REGISTER_OP("CompressElement") 137 .Input("components: input_types") 138 .Output("compressed: variant") 139 .Attr("input_types: list(type) >= 1") 140 .SetShapeFn(shape_inference::ScalarShape); 141 142 REGISTER_OP("UncompressElement") 143 .Input("compressed: variant") 144 .Output("components: output_types") 145 .Attr("output_types: list(type) >= 1") 146 .Attr("output_shapes: list(shape) >= 1") 147 .SetShapeFn(shape_inference::DatasetIteratorShape); 148 149 REGISTER_OP("ComputeBatchSize") 150 .Input("input_dataset : variant") 151 .Output("batch_size : int64") 152 .SetShapeFn(shape_inference::ScalarShape); 153 154 REGISTER_OP("CSVDataset") 155 .Input("filenames: string") 156 .Input("compression_type: string") 157 .Input("buffer_size: int64") 158 .Input("header: bool") 159 .Input("field_delim: string") 160 .Input("use_quote_delim: bool") 161 .Input("na_value: string") 162 .Input("select_cols: int64") 163 .Input("record_defaults: output_types") 164 .Output("handle: variant") 165 .Attr("output_types: list({float,double,int32,int64,string}) >= 1") 166 .Attr("output_shapes: list(shape) >= 1") 167 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 168 // disable constant folding. __anon4afff42d0602(shape_inference::InferenceContext* c) 169 .SetShapeFn([](shape_inference::InferenceContext* c) { 170 shape_inference::ShapeHandle unused; 171 // `filenames` must be a scalar or a vector. 172 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); 173 // `compression_type`, `buffer_size`, `header`, `field_delim`, 174 // `use_quote_delim`, `na_value` must be scalars 175 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 176 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 177 TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); 178 TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused)); 179 TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused)); 180 TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); 181 // `select_cols` must be a vector 182 TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); 183 // `record_defaults` must be lists of scalars 184 for (size_t i = 8; i < c->num_inputs(); ++i) { 185 shape_inference::ShapeHandle v; 186 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); 187 if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { 188 return errors::InvalidArgument( 189 "Shape of a default must be a length-0 or length-1 vector, or a " 190 "scalar."); 191 } 192 } 193 return shape_inference::ScalarShape(c); 194 }); 195 196 REGISTER_OP("CSVDatasetV2") 197 .Input("filenames: string") 198 .Input("compression_type: string") 199 .Input("buffer_size: int64") 200 .Input("header: bool") 201 .Input("field_delim: string") 202 .Input("use_quote_delim: bool") 203 .Input("na_value: string") 204 .Input("select_cols: int64") 205 .Input("record_defaults: output_types") 206 .Input("exclude_cols: int64") 207 .Output("handle: variant") 208 .Attr("output_types: list({float,double,int32,int64,string}) >= 1") 209 .Attr("output_shapes: list(shape) >= 1") 210 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 211 // disable constant folding. __anon4afff42d0702(shape_inference::InferenceContext* c) 212 .SetShapeFn([](shape_inference::InferenceContext* c) { 213 shape_inference::ShapeHandle unused; 214 // `filenames` must be a scalar or a vector. 215 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); 216 // `compression_type`, `buffer_size`, `header`, `field_delim`, 217 // `use_quote_delim`, `na_value` must be scalars 218 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 219 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 220 TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); 221 TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused)); 222 TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused)); 223 TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); 224 // `select_cols` must be a vector 225 TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); 226 // `exclude_cols` must be a vector 227 TF_RETURN_IF_ERROR( 228 c->WithRank(c->input(c->num_inputs() - 1), 1, &unused)); 229 // `record_defaults` must be lists of scalars 230 for (size_t i = 8; i < c->num_inputs() - 1; ++i) { 231 shape_inference::ShapeHandle v; 232 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); 233 if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { 234 return errors::InvalidArgument( 235 "Shape of a default must be a length-0 or length-1 vector, or a " 236 "scalar."); 237 } 238 } 239 return shape_inference::ScalarShape(c); 240 }); 241 242 REGISTER_OP("ExperimentalCSVDataset") 243 .Input("filenames: string") 244 .Input("compression_type: string") 245 .Input("buffer_size: int64") 246 .Input("header: bool") 247 .Input("field_delim: string") 248 .Input("use_quote_delim: bool") 249 .Input("na_value: string") 250 .Input("select_cols: int64") 251 .Input("record_defaults: output_types") 252 .Output("handle: variant") 253 .Attr("output_types: list({float,double,int32,int64,string}) >= 1") 254 .Attr("output_shapes: list(shape) >= 1") 255 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 256 // disable constant folding. __anon4afff42d0802(shape_inference::InferenceContext* c) 257 .SetShapeFn([](shape_inference::InferenceContext* c) { 258 shape_inference::ShapeHandle unused; 259 // `filenames` must be a scalar or a vector. 260 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); 261 // `compression_type`, `buffer_size`, `header`, `field_delim`, 262 // `use_quote_delim`, `na_value` must be scalars 263 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 264 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 265 TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); 266 TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused)); 267 TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused)); 268 TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); 269 // `select_cols` must be a vector 270 TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); 271 // `record_defaults` must be lists of scalars 272 for (size_t i = 8; i < c->num_inputs(); ++i) { 273 shape_inference::ShapeHandle v; 274 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); 275 if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { 276 return errors::InvalidArgument( 277 "Shape of a default must be a length-0 or length-1 vector, or a " 278 "scalar."); 279 } 280 } 281 return shape_inference::ScalarShape(c); 282 }); 283 284 REGISTER_OP("ExperimentalDatasetCardinality") 285 .Input("input_dataset: variant") 286 .Output("cardinality: int64") 287 .SetShapeFn(shape_inference::ScalarShape); 288 289 REGISTER_OP("DatasetFromGraph") 290 .Input("graph_def: string") 291 .Output("handle: variant") 292 .SetShapeFn(shape_inference::ScalarShape); 293 294 // TODO(b/124308596): Instead of conservatively marking this op as stateful, 295 // implement a mechanism to determine whether `dataset` has a side-effect 296 // and use it to decide whether to use a stateless or stateful version of this 297 // op. 298 REGISTER_OP("DatasetToTFRecord") 299 .Input("input_dataset: variant") 300 .Input("filename: string") 301 .Input("compression_type: string") 302 .SetIsStateful() 303 .SetShapeFn(shape_inference::NoOutputs); 304 305 REGISTER_OP("ExperimentalDatasetToTFRecord") 306 .Input("input_dataset: variant") 307 .Input("filename: string") 308 .Input("compression_type: string") 309 .SetIsStateful() 310 .SetShapeFn(shape_inference::NoOutputs); 311 312 REGISTER_OP("DenseToSparseBatchDataset") 313 .Input("input_dataset: variant") 314 .Input("batch_size: int64") 315 .Input("row_shape: int64") 316 .Output("handle: variant") 317 .Attr("output_types: list(type) >= 1") 318 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0902(shape_inference::InferenceContext* c) 319 .SetShapeFn([](shape_inference::InferenceContext* c) { 320 shape_inference::ShapeHandle unused; 321 // batch_size should be a scalar. 322 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 323 // row_shape should be a 1-D vector. 324 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &unused)); 325 return shape_inference::ScalarShape(c); 326 }); 327 328 REGISTER_OP("ExperimentalDenseToSparseBatchDataset") 329 .Input("input_dataset: variant") 330 .Input("batch_size: int64") 331 .Input("row_shape: int64") 332 .Output("handle: variant") 333 .Attr("output_types: list(type) >= 1") 334 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0a02(shape_inference::InferenceContext* c) 335 .SetShapeFn([](shape_inference::InferenceContext* c) { 336 shape_inference::ShapeHandle unused; 337 // batch_size should be a scalar. 338 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 339 // row_shape should be a 1-D vector. 340 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &unused)); 341 return shape_inference::ScalarShape(c); 342 }); 343 344 REGISTER_OP("DirectedInterleaveDataset") 345 .Input("selector_input_dataset: variant") 346 .Input("data_input_datasets: N * variant") 347 .Output("handle: variant") 348 .Attr("output_types: list(type) >= 1") 349 .Attr("output_shapes: list(shape) >= 1") 350 .Attr("N: int >= 1") 351 .SetShapeFn(shape_inference::ScalarShape); 352 353 REGISTER_OP("ExperimentalDirectedInterleaveDataset") 354 .Input("selector_input_dataset: variant") 355 .Input("data_input_datasets: N * variant") 356 .Output("handle: variant") 357 .Attr("output_types: list(type) >= 1") 358 .Attr("output_shapes: list(shape) >= 1") 359 .Attr("N: int >= 1") 360 .SetShapeFn(shape_inference::ScalarShape); 361 362 REGISTER_OP("GroupByReducerDataset") 363 .Input("input_dataset: variant") 364 .Input("key_func_other_arguments: Tkey_func_other_arguments") 365 .Input("init_func_other_arguments: Tinit_func_other_arguments") 366 .Input("reduce_func_other_arguments: Treduce_func_other_arguments") 367 .Input("finalize_func_other_arguments: Tfinalize_func_other_arguments") 368 .Output("handle: variant") 369 .Attr("key_func: func") 370 .Attr("init_func: func") 371 .Attr("reduce_func: func") 372 .Attr("finalize_func: func") 373 .Attr("Tkey_func_other_arguments: list(type) >= 0") 374 .Attr("Tinit_func_other_arguments: list(type) >= 0") 375 .Attr("Treduce_func_other_arguments: list(type) >= 0") 376 .Attr("Tfinalize_func_other_arguments: list(type) >= 0") 377 .Attr("output_types: list(type) >= 1") 378 .Attr("output_shapes: list(shape) >= 1") 379 .SetIsStateful() 380 .SetShapeFn(shape_inference::ScalarShape); 381 382 REGISTER_OP("ExperimentalGroupByReducerDataset") 383 .Input("input_dataset: variant") 384 .Input("key_func_other_arguments: Tkey_func_other_arguments") 385 .Input("init_func_other_arguments: Tinit_func_other_arguments") 386 .Input("reduce_func_other_arguments: Treduce_func_other_arguments") 387 .Input("finalize_func_other_arguments: Tfinalize_func_other_arguments") 388 .Output("handle: variant") 389 .Attr("key_func: func") 390 .Attr("init_func: func") 391 .Attr("reduce_func: func") 392 .Attr("finalize_func: func") 393 .Attr("Tkey_func_other_arguments: list(type) >= 0") 394 .Attr("Tinit_func_other_arguments: list(type) >= 0") 395 .Attr("Treduce_func_other_arguments: list(type) >= 0") 396 .Attr("Tfinalize_func_other_arguments: list(type) >= 0") 397 .Attr("output_types: list(type) >= 1") 398 .Attr("output_shapes: list(shape) >= 1") 399 .SetIsStateful() 400 .SetShapeFn(shape_inference::ScalarShape); 401 402 REGISTER_OP("GroupByWindowDataset") 403 .Input("input_dataset: variant") 404 .Input("key_func_other_arguments: Tkey_func_other_arguments") 405 .Input("reduce_func_other_arguments: Treduce_func_other_arguments") 406 .Input( 407 "window_size_func_other_arguments: Twindow_size_func_other_arguments") 408 .Output("handle: variant") 409 .Attr("key_func: func") 410 .Attr("reduce_func: func") 411 .Attr("window_size_func: func") 412 .Attr("Tkey_func_other_arguments: list(type) >= 0") 413 .Attr("Treduce_func_other_arguments: list(type) >= 0") 414 .Attr("Twindow_size_func_other_arguments: list(type) >= 0") 415 .Attr("output_types: list(type) >= 1") 416 .Attr("output_shapes: list(shape) >= 1") 417 .SetShapeFn(shape_inference::ScalarShape); 418 419 REGISTER_OP("ExperimentalGroupByWindowDataset") 420 .Input("input_dataset: variant") 421 .Input("key_func_other_arguments: Tkey_func_other_arguments") 422 .Input("reduce_func_other_arguments: Treduce_func_other_arguments") 423 .Input( 424 "window_size_func_other_arguments: Twindow_size_func_other_arguments") 425 .Output("handle: variant") 426 .Attr("key_func: func") 427 .Attr("reduce_func: func") 428 .Attr("window_size_func: func") 429 .Attr("Tkey_func_other_arguments: list(type) >= 0") 430 .Attr("Treduce_func_other_arguments: list(type) >= 0") 431 .Attr("Twindow_size_func_other_arguments: list(type) >= 0") 432 .Attr("output_types: list(type) >= 1") 433 .Attr("output_shapes: list(shape) >= 1") 434 .SetShapeFn(shape_inference::ScalarShape); 435 436 REGISTER_OP("IgnoreErrorsDataset") 437 .Input("input_dataset: variant") 438 .Output("handle: variant") 439 .Attr("output_types: list(type) >= 1") 440 .Attr("output_shapes: list(shape) >= 1") 441 .Attr("log_warning: bool = false") 442 .SetShapeFn(shape_inference::ScalarShape); 443 444 REGISTER_OP("ExperimentalIgnoreErrorsDataset") 445 .Input("input_dataset: variant") 446 .Output("handle: variant") 447 .Attr("output_types: list(type) >= 1") 448 .Attr("output_shapes: list(shape) >= 1") 449 .Attr("log_warning: bool = false") 450 .SetShapeFn(shape_inference::ScalarShape); 451 452 REGISTER_OP("IteratorGetDevice") 453 .Input("resource: resource") 454 .Output("device: string") 455 .SetShapeFn(shape_inference::ScalarShape); 456 457 REGISTER_OP("ExperimentalIteratorGetDevice") 458 .Input("resource: resource") 459 .Output("device: string") 460 .SetShapeFn(shape_inference::ScalarShape); 461 462 REGISTER_OP("LatencyStatsDataset") 463 .Input("input_dataset: variant") 464 .Input("tag: string") 465 .Output("handle: variant") 466 .Attr("output_types: list(type) >= 1") 467 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0b02(shape_inference::InferenceContext* c) 468 .SetShapeFn([](shape_inference::InferenceContext* c) { 469 shape_inference::ShapeHandle tag_shape; 470 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &tag_shape)); 471 return shape_inference::ScalarShape(c); 472 }); 473 474 REGISTER_OP("ExperimentalLatencyStatsDataset") 475 .Input("input_dataset: variant") 476 .Input("tag: string") 477 .Output("handle: variant") 478 .Attr("output_types: list(type) >= 1") 479 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d0c02(shape_inference::InferenceContext* c) 480 .SetShapeFn([](shape_inference::InferenceContext* c) { 481 shape_inference::ShapeHandle tag_shape; 482 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &tag_shape)); 483 return shape_inference::ScalarShape(c); 484 }); 485 486 REGISTER_OP("LMDBDataset") 487 .Input("filenames: string") 488 .Output("handle: variant") 489 .Attr("output_types: list(type) >= 1") 490 .Attr("output_shapes: list(shape) >= 1") 491 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 492 // disable constant folding. 493 .SetShapeFn(shape_inference::ScalarShape); 494 495 REGISTER_OP("ExperimentalLMDBDataset") 496 .Input("filenames: string") 497 .Output("handle: variant") 498 .Attr("output_types: list(type) >= 1") 499 .Attr("output_shapes: list(shape) >= 1") 500 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 501 // disable constant folding. 502 .SetShapeFn(shape_inference::ScalarShape); 503 504 REGISTER_OP("MapAndBatchDataset") 505 .Input("input_dataset: variant") 506 .Input("other_arguments: Targuments") 507 .Input("batch_size: int64") 508 .Input("num_parallel_calls: int64") 509 .Input("drop_remainder: bool") 510 .Output("handle: variant") 511 .Attr("f: func") 512 .Attr("Targuments: list(type) >= 0") 513 .Attr("output_types: list(type) >= 1") 514 .Attr("output_shapes: list(shape) >= 1") 515 .Attr("preserve_cardinality: bool = false") __anon4afff42d0d02(shape_inference::InferenceContext* c) 516 .SetShapeFn([](shape_inference::InferenceContext* c) { 517 // Use index from the end to retrieve the Input shapes, 518 // so that to avoid guessing the length of "other_arguments". 519 // batch_size, num_parallel_calls, and drop_remainder are 0-D scalars. 520 shape_inference::ShapeHandle unused; 521 TF_RETURN_IF_ERROR( 522 c->WithRank(c->input(c->num_inputs() - 3), 0, &unused)); 523 TF_RETURN_IF_ERROR( 524 c->WithRank(c->input(c->num_inputs() - 2), 0, &unused)); 525 TF_RETURN_IF_ERROR( 526 c->WithRank(c->input(c->num_inputs() - 1), 0, &unused)); 527 528 return shape_inference::ScalarShape(c); 529 }); 530 531 REGISTER_OP("ExperimentalMapAndBatchDataset") 532 .Input("input_dataset: variant") 533 .Input("other_arguments: Targuments") 534 .Input("batch_size: int64") 535 .Input("num_parallel_calls: int64") 536 .Input("drop_remainder: bool") 537 .Output("handle: variant") 538 .Attr("f: func") 539 .Attr("Targuments: list(type) >= 0") 540 .Attr("output_types: list(type) >= 1") 541 .Attr("output_shapes: list(shape) >= 1") 542 .Attr("preserve_cardinality: bool = false") __anon4afff42d0e02(shape_inference::InferenceContext* c) 543 .SetShapeFn([](shape_inference::InferenceContext* c) { 544 // Use index from the end to retrieve the Input shapes, 545 // so that to avoid guessing the length of "other_arguments". 546 // batch_size, num_parallel_calls, and drop_remainder are 0-D scalars. 547 shape_inference::ShapeHandle unused; 548 TF_RETURN_IF_ERROR( 549 c->WithRank(c->input(c->num_inputs() - 3), 0, &unused)); 550 TF_RETURN_IF_ERROR( 551 c->WithRank(c->input(c->num_inputs() - 2), 0, &unused)); 552 TF_RETURN_IF_ERROR( 553 c->WithRank(c->input(c->num_inputs() - 1), 0, &unused)); 554 555 return shape_inference::ScalarShape(c); 556 }); 557 558 REGISTER_OP("ExperimentalMapDataset") 559 .Input("input_dataset: variant") 560 .Input("other_arguments: Targuments") 561 .Output("handle: variant") 562 .Attr("f: func") 563 .Attr("Targuments: list(type) >= 0") 564 .Attr("output_types: list(type) >= 1") 565 .Attr("output_shapes: list(shape) >= 1") 566 .Attr("use_inter_op_parallelism: bool = true") 567 .Attr("preserve_cardinality: bool = false") 568 .SetShapeFn(shape_inference::ScalarShape); 569 570 REGISTER_OP("MatchingFilesDataset") 571 .Input("patterns: string") 572 .Output("handle: variant") 573 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 574 // disable constant folding. __anon4afff42d0f02(shape_inference::InferenceContext* c) 575 .SetShapeFn([](shape_inference::InferenceContext* c) { 576 shape_inference::ShapeHandle unused; 577 // `patterns` must be a scalar or a vector. 578 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); 579 return shape_inference::ScalarShape(c); 580 }); 581 582 REGISTER_OP("ExperimentalMatchingFilesDataset") 583 .Input("patterns: string") 584 .Output("handle: variant") 585 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 586 // disable constant folding. __anon4afff42d1002(shape_inference::InferenceContext* c) 587 .SetShapeFn([](shape_inference::InferenceContext* c) { 588 shape_inference::ShapeHandle unused; 589 // `patterns` must be a scalar or a vector. 590 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); 591 return shape_inference::ScalarShape(c); 592 }); 593 594 REGISTER_OP("MaxIntraOpParallelismDataset") 595 .Input("input_dataset: variant") 596 .Input("max_intra_op_parallelism: int64") 597 .Output("handle: variant") 598 .Attr("output_types: list(type) >= 1") 599 .Attr("output_shapes: list(shape) >= 1") 600 .SetShapeFn(shape_inference::ScalarShape); 601 602 REGISTER_OP("ExperimentalMaxIntraOpParallelismDataset") 603 .Input("input_dataset: variant") 604 .Input("max_intra_op_parallelism: int64") 605 .Output("handle: variant") 606 .Attr("output_types: list(type) >= 1") 607 .Attr("output_shapes: list(shape) >= 1") 608 .SetShapeFn(shape_inference::ScalarShape); 609 610 REGISTER_OP("NonSerializableDataset") 611 .Input("input_dataset: variant") 612 .Output("handle: variant") 613 .Attr("output_types: list(type) >= 1") 614 .Attr("output_shapes: list(shape) >= 1") 615 .SetShapeFn(shape_inference::ScalarShape); 616 617 REGISTER_OP("ExperimentalNonSerializableDataset") 618 .Input("input_dataset: variant") 619 .Output("handle: variant") 620 .Attr("output_types: list(type) >= 1") 621 .Attr("output_shapes: list(shape) >= 1") 622 .SetShapeFn(shape_inference::ScalarShape); 623 624 REGISTER_OP("ParallelInterleaveDataset") 625 .Input("input_dataset: variant") 626 .Input("other_arguments: Targuments") 627 .Input("cycle_length: int64") 628 .Input("block_length: int64") 629 .Input("sloppy: bool") 630 .Input("buffer_output_elements: int64") 631 .Input("prefetch_input_elements: int64") 632 .Output("handle: variant") 633 .Attr("f: func") 634 .Attr("Targuments: list(type) >= 0") 635 .Attr("output_types: list(type) >= 1") 636 .Attr("output_shapes: list(shape) >= 1") 637 .SetShapeFn(shape_inference::ScalarShape); 638 639 // This is the V2 of ParallelInterleaveDataset, renamed to differentiate it 640 // from the non-experimental ParallelInterleaveDataset op. 641 REGISTER_OP("LegacyParallelInterleaveDatasetV2") 642 .Input("input_dataset: variant") 643 .Input("other_arguments: Targuments") 644 .Input("cycle_length: int64") 645 .Input("block_length: int64") 646 .Input("buffer_output_elements: int64") 647 .Input("prefetch_input_elements: int64") 648 .Output("handle: variant") 649 .Attr("f: func") 650 // "true", "false", or "default". 651 .Attr("deterministic: string = 'default'") 652 .Attr("Targuments: list(type) >= 0") 653 .Attr("output_types: list(type) >= 1") 654 .Attr("output_shapes: list(shape) >= 1") 655 .SetShapeFn(shape_inference::ScalarShape); 656 657 // This op is no longer used. We keep it so that we can read graphs written by 658 // old versions of TensorFlow. 659 REGISTER_OP("ExperimentalParallelInterleaveDataset") 660 .Input("input_dataset: variant") 661 .Input("other_arguments: Targuments") 662 .Input("cycle_length: int64") 663 .Input("block_length: int64") 664 .Input("sloppy: bool") 665 .Input("buffer_output_elements: int64") 666 .Input("prefetch_input_elements: int64") 667 .Output("handle: variant") 668 .Attr("f: func") 669 .Attr("Targuments: list(type) >= 0") 670 .Attr("output_types: list(type) >= 1") 671 .Attr("output_shapes: list(shape) >= 1") 672 .SetShapeFn(shape_inference::ScalarShape); 673 674 REGISTER_OP("ParseExampleDataset") 675 .Input("input_dataset: variant") 676 .Input("num_parallel_calls: int64") 677 .Input("dense_defaults: Tdense") 678 .Output("handle: variant") 679 .Attr("sparse_keys: list(string) >= 0") 680 .Attr("dense_keys: list(string) >= 0") 681 .Attr("sparse_types: list({float,int64,string}) >= 0") 682 .Attr("Tdense: list({float,int64,string}) >= 0") 683 .Attr("dense_shapes: list(shape) >= 0") 684 .Attr("output_types: list(type) >= 1") 685 .Attr("output_shapes: list(shape) >= 1") // Output components will be 686 // sorted by key (dense_keys and 687 // sparse_keys combined) here. 688 .Attr("sloppy: bool = false") 689 .Attr("ragged_keys: list(string) >= 0 = []") 690 .Attr("ragged_value_types: list({float,int64,string}) >= 0 = []") 691 .Attr("ragged_split_types: list({int32,int64}) >= 0 = []") 692 .SetShapeFn(shape_inference::ScalarShape); 693 694 REGISTER_OP("ParseExampleDatasetV2") 695 .Input("input_dataset: variant") 696 .Input("num_parallel_calls: int64") 697 .Input("dense_defaults: Tdense") 698 .Output("handle: variant") 699 .Attr("sparse_keys: list(string) >= 0") 700 .Attr("dense_keys: list(string) >= 0") 701 .Attr("sparse_types: list({float,int64,string}) >= 0") 702 .Attr("Tdense: list({float,int64,string}) >= 0") 703 .Attr("dense_shapes: list(shape) >= 0") 704 .Attr("output_types: list(type) >= 1") 705 .Attr("output_shapes: list(shape) >= 1") // Output components will be 706 // sorted by key (dense_keys and 707 // sparse_keys combined) here. 708 // "true", "false", or "default". 709 .Attr("deterministic: string = 'default'") 710 .Attr("ragged_keys: list(string) >= 0 = []") 711 .Attr("ragged_value_types: list({float,int64,string}) >= 0 = []") 712 .Attr("ragged_split_types: list({int32,int64}) >= 0 = []") 713 .SetShapeFn(shape_inference::ScalarShape); 714 715 REGISTER_OP("ExperimentalParseExampleDataset") 716 .Input("input_dataset: variant") 717 .Input("num_parallel_calls: int64") 718 .Input("dense_defaults: Tdense") 719 .Output("handle: variant") 720 .Attr("sparse_keys: list(string) >= 0") 721 .Attr("dense_keys: list(string) >= 0") 722 .Attr("sparse_types: list({float,int64,string}) >= 0") 723 .Attr("Tdense: list({float,int64,string}) >= 0") 724 .Attr("dense_shapes: list(shape) >= 0") 725 .Attr("output_types: list(type) >= 1") 726 .Attr("output_shapes: list(shape) >= 1") // Output components will be 727 // sorted by key (dense_keys and 728 // sparse_keys combined) here. 729 .Attr("sloppy: bool = false") 730 .SetShapeFn(shape_inference::ScalarShape); 731 732 REGISTER_OP("PrivateThreadPoolDataset") 733 .Input("input_dataset: variant") 734 .Input("num_threads: int64") 735 .Output("handle: variant") 736 .Attr("output_types: list(type) >= 1") 737 .Attr("output_shapes: list(shape) >= 1") 738 .SetShapeFn(shape_inference::ScalarShape); 739 740 REGISTER_OP("ExperimentalPrivateThreadPoolDataset") 741 .Input("input_dataset: variant") 742 .Input("num_threads: int64") 743 .Output("handle: variant") 744 .Attr("output_types: list(type) >= 1") 745 .Attr("output_shapes: list(shape) >= 1") 746 .SetShapeFn(shape_inference::ScalarShape); 747 748 REGISTER_OP("ExperimentalRandomDataset") 749 .Input("seed: int64") 750 .Input("seed2: int64") 751 .Output("handle: variant") 752 .Attr("output_types: list(type) >= 1") 753 .Attr("output_shapes: list(shape) >= 1") 754 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 755 // disable constant folding. __anon4afff42d1102(shape_inference::InferenceContext* c) 756 .SetShapeFn([](shape_inference::InferenceContext* c) { 757 shape_inference::ShapeHandle unused; 758 // buffer_size, seed, and seed2 should be scalars. 759 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused)); 760 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 761 return shape_inference::ScalarShape(c); 762 }); 763 764 REGISTER_OP("RandomDataset") 765 .Input("seed: int64") 766 .Input("seed2: int64") 767 .Output("handle: variant") 768 .Attr("output_types: list(type) >= 1") 769 .Attr("output_shapes: list(shape) >= 1") 770 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 771 // disable constant folding. __anon4afff42d1202(shape_inference::InferenceContext* c) 772 .SetShapeFn([](shape_inference::InferenceContext* c) { 773 shape_inference::ShapeHandle unused; 774 // buffer_size, seed, and seed2 should be scalars. 775 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused)); 776 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 777 return shape_inference::ScalarShape(c); 778 }); 779 780 REGISTER_OP("ExperimentalRebatchDataset") 781 .Input("input_dataset: variant") 782 .Input("num_replicas: int64") 783 .Output("handle: variant") 784 .Attr("output_types: list(type) >= 1") 785 .Attr("output_shapes: list(shape) >= 1") 786 .Attr("use_fallback: bool = true") 787 .SetShapeFn(shape_inference::ScalarShape); 788 789 REGISTER_OP("RebatchDataset") 790 .Input("input_dataset: variant") 791 .Input("num_replicas: int64") 792 .Output("handle: variant") 793 .Attr("output_types: list(type) >= 1") 794 .Attr("output_shapes: list(shape) >= 1") 795 .Attr("use_fallback: bool = true") 796 .SetShapeFn(shape_inference::ScalarShape); 797 798 REGISTER_OP("RebatchDatasetV2") 799 .Input("input_dataset: variant") 800 .Input("batch_sizes: int64") 801 .Input("drop_remainder: bool") 802 .Output("handle: variant") 803 .Attr("output_types: list(type) >= 1") 804 .Attr("output_shapes: list(shape) >= 1") 805 .SetShapeFn(shape_inference::ScalarShape); 806 807 REGISTER_OP("SamplingDataset") 808 .Input("input_dataset: variant") 809 .Input("rate: float32") 810 .Input("seed: int64") 811 .Input("seed2: int64") 812 .Output("handle: variant") 813 .Attr("output_types: list(type) >= 1") 814 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d1302(shape_inference::InferenceContext* c) 815 .SetShapeFn([](shape_inference::InferenceContext* c) { 816 shape_inference::ShapeHandle unused; 817 // rate, seed, and seed2 should be scalars. 818 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 819 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 820 TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); 821 return shape_inference::ScalarShape(c); 822 }); 823 824 REGISTER_OP("ScanDataset") 825 .Input("input_dataset: variant") 826 .Input("initial_state: Tstate") 827 .Input("other_arguments: Targuments") 828 .Output("handle: variant") 829 .Attr("f: func") 830 .Attr("Tstate: list(type) >= 1") 831 .Attr("Targuments: list(type) >= 0") 832 .Attr("output_types: list(type) >= 1") 833 .Attr("output_shapes: list(shape) >= 1") 834 .Attr("preserve_cardinality: bool = false") 835 .Attr("use_default_device: bool = true") 836 .SetShapeFn(shape_inference::ScalarShape); 837 838 REGISTER_OP("ExperimentalScanDataset") 839 .Input("input_dataset: variant") 840 .Input("initial_state: Tstate") 841 .Input("other_arguments: Targuments") 842 .Output("handle: variant") 843 .Attr("f: func") 844 .Attr("Tstate: list(type) >= 1") 845 .Attr("Targuments: list(type) >= 0") 846 .Attr("output_types: list(type) >= 1") 847 .Attr("output_shapes: list(shape) >= 1") 848 .Attr("preserve_cardinality: bool = false") 849 .SetShapeFn(shape_inference::ScalarShape); 850 851 REGISTER_OP("SetStatsAggregatorDataset") 852 .Input("input_dataset: variant") 853 .Input("stats_aggregator: resource") 854 .Input("tag: string") 855 .Input("counter_prefix: string") 856 .Output("handle: variant") 857 .Attr("output_types: list(type) >= 1") 858 .Attr("output_shapes: list(shape) >= 1") 859 .SetShapeFn(shape_inference::ScalarShape); 860 861 REGISTER_OP("ExperimentalSetStatsAggregatorDataset") 862 .Input("input_dataset: variant") 863 .Input("stats_aggregator: resource") 864 .Input("tag: string") 865 .Input("counter_prefix: string") 866 .Output("handle: variant") 867 .Attr("output_types: list(type) >= 1") 868 .Attr("output_shapes: list(shape) >= 1") 869 .SetShapeFn(shape_inference::ScalarShape); 870 871 REGISTER_OP("SleepDataset") 872 .Input("input_dataset: variant") 873 .Input("sleep_microseconds: int64") 874 .Output("handle: variant") 875 .Attr("output_types: list(type) >= 1") 876 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d1402(shape_inference::InferenceContext* c) 877 .SetShapeFn([](shape_inference::InferenceContext* c) { 878 shape_inference::ShapeHandle unused; 879 // Both inputs are scalar. 880 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 0, &unused)); 881 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 0, &unused)); 882 return shape_inference::ScalarShape(c); 883 }); 884 885 REGISTER_OP("ExperimentalSleepDataset") 886 .Input("input_dataset: variant") 887 .Input("sleep_microseconds: int64") 888 .Output("handle: variant") 889 .Attr("output_types: list(type) >= 1") 890 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d1502(shape_inference::InferenceContext* c) 891 .SetShapeFn([](shape_inference::InferenceContext* c) { 892 shape_inference::ShapeHandle unused; 893 // Both inputs are scalar. 894 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 0, &unused)); 895 TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 0, &unused)); 896 return shape_inference::ScalarShape(c); 897 }); 898 899 REGISTER_OP("SlidingWindowDataset") 900 .Input("input_dataset: variant") 901 .Input("window_size: int64") 902 .Input("window_shift: int64") 903 .Input("window_stride: int64") 904 .Output("handle: variant") 905 .Attr("output_types: list(type) >= 1") 906 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d1602(shape_inference::InferenceContext* c) 907 .SetShapeFn([](shape_inference::InferenceContext* c) { 908 shape_inference::ShapeHandle unused; 909 // window_size, window_shift, and window_stride should be scalars. 910 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 911 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 912 TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); 913 return shape_inference::ScalarShape(c); 914 }); 915 916 REGISTER_OP("ExperimentalSlidingWindowDataset") 917 .Input("input_dataset: variant") 918 .Input("window_size: int64") 919 .Input("window_shift: int64") 920 .Input("window_stride: int64") 921 .Output("handle: variant") 922 .Attr("output_types: list(type) >= 1") 923 .Attr("output_shapes: list(shape) >= 1") __anon4afff42d1702(shape_inference::InferenceContext* c) 924 .SetShapeFn([](shape_inference::InferenceContext* c) { 925 shape_inference::ShapeHandle unused; 926 // window_size, window_shift, and window_stride should be scalars. 927 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 928 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 929 TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); 930 return shape_inference::ScalarShape(c); 931 }); 932 933 REGISTER_OP("SnapshotDataset") 934 .Input("input_dataset: variant") 935 .Input("path: string") 936 .Output("handle: variant") 937 .Attr("output_types: list(type) >= 1") 938 .Attr("output_shapes: list(shape) >= 1") 939 .Attr("compression: string = ''") 940 .Attr("reader_path_prefix: string = ''") 941 .Attr("writer_path_prefix: string = ''") 942 .Attr("shard_size_bytes: int = 10737418240") // 10 GiB default 943 .Attr("pending_snapshot_expiry_seconds: int = 86400") // 1 day default 944 .Attr("num_reader_threads: int = 1") 945 .Attr("reader_buffer_size: int = 1") 946 .Attr("num_writer_threads: int = 1") 947 .Attr("writer_buffer_size: int = 1") 948 .Attr("shuffle_on_read: bool = false") 949 .Attr("seed: int = 0") 950 .Attr("seed2: int = 0") 951 .Attr("mode: string = 'auto'") 952 .Attr("snapshot_name: string = ''") __anon4afff42d1802(shape_inference::InferenceContext* c) 953 .SetShapeFn([](shape_inference::InferenceContext* c) { 954 shape_inference::ShapeHandle unused; 955 // snapshot_path should be a scalar. 956 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 957 return shape_inference::ScalarShape(c); 958 }); 959 960 REGISTER_OP("SnapshotDatasetV2") 961 .Input("input_dataset: variant") 962 .Input("path: string") 963 .Input("reader_func_other_args: Treader_func_args") 964 .Input("shard_func_other_args: Tshard_func_args") 965 .Output("handle: variant") 966 .Attr("output_types: list(type) >= 1") 967 .Attr("output_shapes: list(shape) >= 1") 968 .Attr("compression: string = ''") 969 .Attr("reader_prefix: string = ''") 970 .Attr("writer_prefix: string = ''") 971 .Attr("hash_valid: bool = false") 972 .Attr("hash: int = 0") 973 .Attr("reader_func: func") 974 .Attr("shard_func: func") 975 .Attr("Treader_func_args: list(type) >= 0") 976 .Attr("Tshard_func_args: list(type) >= 0") __anon4afff42d1902(shape_inference::InferenceContext* c) 977 .SetShapeFn([](shape_inference::InferenceContext* c) { 978 shape_inference::ShapeHandle unused; 979 // `path` should be a scalar. 980 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 981 return shape_inference::ScalarShape(c); 982 }); 983 984 REGISTER_OP("SaveDataset") 985 .Input("input_dataset: variant") 986 .Input("path: string") 987 .Input("shard_func_other_args: Tshard_func_args") 988 .Attr("compression: string = ''") 989 .Attr("shard_func: func") 990 .Attr("use_shard_func: bool = true") 991 .Attr("Tshard_func_args: list(type) >= 0") 992 .SetIsStateful() __anon4afff42d1a02(shape_inference::InferenceContext* c) 993 .SetShapeFn([](shape_inference::InferenceContext* c) { 994 shape_inference::ShapeHandle unused; 995 // `path` should be a scalar. 996 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 997 return Status::OK(); 998 }); 999 1000 REGISTER_OP("LoadDataset") 1001 .Input("path: string") 1002 .Input("reader_func_other_args: Treader_func_args") 1003 .Output("handle: variant") 1004 .Attr("output_types: list(type) >= 1") 1005 .Attr("output_shapes: list(shape) >= 1") 1006 .Attr("compression: string = ''") 1007 .Attr("reader_func: func") 1008 .Attr("Treader_func_args: list(type) >= 0") 1009 .SetIsStateful() __anon4afff42d1b02(shape_inference::InferenceContext* c) 1010 .SetShapeFn([](shape_inference::InferenceContext* c) { 1011 shape_inference::ShapeHandle unused; 1012 // `path` should be a scalar. 1013 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused)); 1014 return shape_inference::ScalarShape(c); 1015 }); 1016 1017 REGISTER_OP("SqlDataset") 1018 .Input("driver_name: string") 1019 .Input("data_source_name: string") 1020 .Input("query: string") 1021 .Output("handle: variant") 1022 .Attr("output_types: list(type) >= 1") 1023 .Attr("output_shapes: list(shape) >= 1") 1024 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 1025 // disable constant folding. __anon4afff42d1c02(shape_inference::InferenceContext* c) 1026 .SetShapeFn([](shape_inference::InferenceContext* c) { 1027 shape_inference::ShapeHandle unused; 1028 // driver_name, data_source_name, and query should be scalars. 1029 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused)); 1030 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 1031 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 1032 return shape_inference::ScalarShape(c); 1033 }); 1034 1035 REGISTER_OP("ExperimentalSqlDataset") 1036 .Input("driver_name: string") 1037 .Input("data_source_name: string") 1038 .Input("query: string") 1039 .Output("handle: variant") 1040 .Attr("output_types: list(type) >= 1") 1041 .Attr("output_shapes: list(shape) >= 1") 1042 .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must 1043 // disable constant folding. __anon4afff42d1d02(shape_inference::InferenceContext* c) 1044 .SetShapeFn([](shape_inference::InferenceContext* c) { 1045 shape_inference::ShapeHandle unused; 1046 // driver_name, data_source_name, and query should be scalars. 1047 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused)); 1048 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); 1049 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); 1050 return shape_inference::ScalarShape(c); 1051 }); 1052 1053 REGISTER_OP("StatsAggregatorHandle") 1054 .Output("handle: resource") 1055 .SetShapeFn(shape_inference::ScalarShape) 1056 .Attr("container: string = ''") 1057 .Attr("shared_name: string = ''"); 1058 1059 REGISTER_OP("ExperimentalStatsAggregatorHandle") 1060 .Output("handle: resource") 1061 .SetShapeFn(shape_inference::ScalarShape) 1062 .Attr("container: string = ''") 1063 .Attr("shared_name: string = ''"); 1064 1065 REGISTER_OP("StatsAggregatorHandleV2") 1066 .Output("handle: resource") 1067 .SetShapeFn(shape_inference::ScalarShape) 1068 .Attr("container: string = ''") 1069 .Attr("shared_name: string = ''"); 1070 1071 REGISTER_OP("StatsAggregatorSetSummaryWriter") 1072 .Input("stats_aggregator: resource") 1073 .Input("summary: resource") 1074 .SetShapeFn(shape_inference::NoOutputs); 1075 1076 REGISTER_OP("StatsAggregatorSummary") 1077 .Input("iterator: resource") 1078 .Output("summary: string") 1079 .SetShapeFn(shape_inference::ScalarShape); 1080 1081 REGISTER_OP("ExperimentalStatsAggregatorSummary") 1082 .Input("iterator: resource") 1083 .Output("summary: string") 1084 .SetShapeFn(shape_inference::ScalarShape); 1085 1086 REGISTER_OP("TakeWhileDataset") 1087 .Input("input_dataset: variant") 1088 .Input("other_arguments: Targuments") 1089 .Output("handle: variant") 1090 .Attr("predicate: func") 1091 .Attr("Targuments: list(type) >= 0") 1092 .Attr("output_types: list(type) >= 1") 1093 .Attr("output_shapes: list(shape) >= 1") 1094 .SetShapeFn(shape_inference::ScalarShape); 1095 1096 REGISTER_OP("ExperimentalTakeWhileDataset") 1097 .Input("input_dataset: variant") 1098 .Input("other_arguments: Targuments") 1099 .Output("handle: variant") 1100 .Attr("predicate: func") 1101 .Attr("Targuments: list(type) >= 0") 1102 .Attr("output_types: list(type) >= 1") 1103 .Attr("output_shapes: list(shape) >= 1") 1104 .SetShapeFn(shape_inference::ScalarShape); 1105 1106 REGISTER_OP("ThreadPoolDataset") 1107 .Input("input_dataset: variant") 1108 .Input("thread_pool: resource") 1109 .Output("handle: variant") 1110 .Attr("output_types: list(type) >= 1") 1111 .Attr("output_shapes: list(shape) >= 1") 1112 .SetShapeFn(shape_inference::ScalarShape); 1113 1114 REGISTER_OP("ExperimentalThreadPoolDataset") 1115 .Input("input_dataset: variant") 1116 .Input("thread_pool: resource") 1117 .Output("handle: variant") 1118 .Attr("output_types: list(type) >= 1") 1119 .Attr("output_shapes: list(shape) >= 1") 1120 .SetShapeFn(shape_inference::ScalarShape); 1121 1122 REGISTER_OP("ThreadPoolHandle") 1123 .Output("handle: resource") 1124 .SetShapeFn(shape_inference::ScalarShape) 1125 .Attr("num_threads: int") 1126 .Attr("max_intra_op_parallelism: int = 1") 1127 .Attr("display_name: string") 1128 .Attr("container: string = ''") 1129 .Attr("shared_name: string = ''"); 1130 1131 REGISTER_OP("ExperimentalThreadPoolHandle") 1132 .Output("handle: resource") 1133 .SetShapeFn(shape_inference::ScalarShape) 1134 .Attr("num_threads: int") 1135 .Attr("max_intra_op_parallelism: int = 1") 1136 .Attr("display_name: string") 1137 .Attr("container: string = ''") 1138 .Attr("shared_name: string = ''"); 1139 1140 REGISTER_OP("UnbatchDataset") 1141 .Input("input_dataset: variant") 1142 .Output("handle: variant") 1143 .Attr("output_types: list(type) >= 1") 1144 .Attr("output_shapes: list(shape) >= 1") 1145 .SetShapeFn(shape_inference::ScalarShape); 1146 1147 REGISTER_OP("ExperimentalUnbatchDataset") 1148 .Input("input_dataset: variant") 1149 .Output("handle: variant") 1150 .Attr("output_types: list(type) >= 1") 1151 .Attr("output_shapes: list(shape) >= 1") 1152 .SetShapeFn(shape_inference::ScalarShape); 1153 1154 REGISTER_OP("UniqueDataset") 1155 .Input("input_dataset: variant") 1156 .Output("handle: variant") 1157 .Attr("output_types: list(type) >= 1") 1158 .Attr("output_shapes: list(shape) >= 1") 1159 .SetShapeFn(shape_inference::ScalarShape); 1160 1161 REGISTER_OP("ExperimentalUniqueDataset") 1162 .Input("input_dataset: variant") 1163 .Output("handle: variant") 1164 .Attr("output_types: list(type) >= 1") 1165 .Attr("output_shapes: list(shape) >= 1") 1166 .SetShapeFn(shape_inference::ScalarShape); 1167 1168 REGISTER_OP("DummyIterationCounter") 1169 .Output("handle: resource") __anon4afff42d1e02(shape_inference::InferenceContext* c) 1170 .SetShapeFn([](shape_inference::InferenceContext* c) { 1171 c->set_output(0, c->Scalar()); 1172 return Status::OK(); 1173 }); 1174 1175 REGISTER_OP("DataServiceDataset") 1176 .Input("dataset_id: int64") 1177 .Input("processing_mode: string") 1178 .Input("address: string") 1179 .Input("protocol: string") 1180 .Input("job_name: string") 1181 .Input("max_outstanding_requests: int64") 1182 .Input("iteration_counter: resource") 1183 .Output("handle: variant") 1184 .Attr("task_refresh_interval_hint_ms: int = -1") 1185 .Attr("output_types: list(type) >= 1") 1186 .Attr("output_shapes: list(shape) >= 1") 1187 .Attr("data_transfer_protocol: string = ''") 1188 .SetIsStateful() 1189 .SetShapeFn(shape_inference::ScalarShape); 1190 1191 // Adds `consumer_index` and `num_consumers` arguments to support round-robin 1192 // reads. 1193 REGISTER_OP("DataServiceDatasetV2") 1194 .Input("dataset_id: int64") 1195 .Input("processing_mode: string") 1196 .Input("address: string") 1197 .Input("protocol: string") 1198 .Input("job_name: string") 1199 .Input("consumer_index: int64") 1200 .Input("num_consumers: int64") 1201 .Input("max_outstanding_requests: int64") 1202 .Input("iteration_counter: resource") 1203 .Output("handle: variant") 1204 .Attr("task_refresh_interval_hint_ms: int = -1") 1205 .Attr("output_types: list(type) >= 1") 1206 .Attr("output_shapes: list(shape) >= 1") 1207 .Attr("data_transfer_protocol: string = ''") 1208 .SetIsStateful() 1209 .SetShapeFn(shape_inference::ScalarShape); 1210 1211 REGISTER_OP("RegisterDataset") 1212 .Input("dataset: variant") 1213 .Input("address: string") 1214 .Input("protocol: string") 1215 .Output("dataset_id: int64") 1216 .Attr("external_state_policy: int") 1217 .SetShapeFn(shape_inference::ScalarShape); 1218 1219 } // namespace tensorflow 1220