1 | /* Copyright 2015 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 | |
16 | #include "tensorflow/core/framework/common_shape_fns.h" |
17 | #include "tensorflow/core/framework/dataset_stateful_op_allowlist.h" |
18 | #include "tensorflow/core/framework/op.h" |
19 | #include "tensorflow/core/framework/shape_inference.h" |
20 | |
21 | namespace tensorflow { |
22 | |
23 | using shape_inference::InferenceContext; |
24 | |
25 | REGISTER_OP("Assert" ) |
26 | .Input("condition: bool" ) |
27 | .Input("data: T" ) |
28 | .SetIsStateful() |
29 | .Attr("T: list(type)" ) |
30 | .Attr("summarize: int = 3" ) |
31 | .SetShapeFn(shape_inference::NoOutputs); |
32 | |
33 | ALLOW_STATEFUL_OP_FOR_DATASET_FUNCTIONS("Assert" ); |
34 | |
35 | REGISTER_OP("Print" ) |
36 | .Input("input: T" ) |
37 | .Input("data: U" ) |
38 | .Output("output: T" ) |
39 | .SetIsStateful() |
40 | .Attr("T: type" ) |
41 | .Attr("U: list(type) >= 0" ) |
42 | .Attr("message: string = ''" ) |
43 | .Attr("first_n: int = -1" ) |
44 | .Attr("summarize: int = 3" ) |
45 | .SetShapeFn(shape_inference::UnchangedShape); |
46 | |
47 | ALLOW_STATEFUL_OP_FOR_DATASET_FUNCTIONS("Print" ); |
48 | |
49 | REGISTER_OP("PrintV2" ) |
50 | .Input("input: string" ) |
51 | .SetIsStateful() |
52 | .Attr("output_stream: string = 'stderr'" ) |
53 | .Attr("end: string = '\n'" ) |
54 | .SetShapeFn([](InferenceContext* c) { |
55 | // Early exit if rank is unknown. |
56 | if (!c->RankKnown(c->input(0))) return OkStatus(); |
57 | // Make sure that the input is a scalar. |
58 | if (c->Rank(c->input(0)) != 0) { |
59 | return errors::InvalidArgument("input must be a scalar, but has rank: " , |
60 | c->Rank(c->input(0))); |
61 | } |
62 | return OkStatus(); |
63 | }); |
64 | |
65 | ALLOW_STATEFUL_OP_FOR_DATASET_FUNCTIONS("PrintV2" ); |
66 | |
67 | // ---------------------------------------------------------------------------- |
68 | // Operators that deal with SummaryProtos (encoded as DT_STRING tensors) as |
69 | // inputs or outputs in various ways. |
70 | |
71 | REGISTER_OP("TensorSummaryV2" ) |
72 | .Input("tag: string" ) |
73 | .Input("tensor: T" ) |
74 | // This serialized summary metadata field describes a summary value, |
75 | // specifically which plugins may use that summary. |
76 | .Input("serialized_summary_metadata: string" ) |
77 | .Output("summary: string" ) |
78 | .Attr("T: type" ) |
79 | .SetShapeFn(shape_inference::ScalarShape); |
80 | |
81 | REGISTER_OP("TensorSummary" ) |
82 | .Input("tensor: T" ) |
83 | .Output("summary: string" ) |
84 | .Attr("T: type" ) |
85 | .Attr("description: string = ''" ) |
86 | .Attr("labels: list(string) = []" ) |
87 | .Attr("display_name: string = ''" ) |
88 | .SetShapeFn(shape_inference::ScalarShape); |
89 | |
90 | REGISTER_OP("ImageSummary" ) |
91 | .Input("tag: string" ) |
92 | .Input("tensor: T" ) |
93 | .Output("summary: string" ) |
94 | .Attr("max_images: int >= 1 = 3" ) |
95 | .Attr("T: {uint8, float, half, float64} = DT_FLOAT" ) |
96 | .Attr( |
97 | "bad_color: tensor = { dtype: DT_UINT8 " |
98 | "tensor_shape: { dim { size: 4 } } " |
99 | "int_val: 255 int_val: 0 int_val: 0 int_val: 255 }" ) |
100 | .SetShapeFn(shape_inference::ScalarShape); |
101 | |
102 | REGISTER_OP("AudioSummaryV2" ) |
103 | .Input("tag: string" ) |
104 | .Input("tensor: float" ) |
105 | .Input("sample_rate: float" ) |
106 | .Output("summary: string" ) |
107 | .Attr("max_outputs: int >= 1 = 3" ) |
108 | .SetShapeFn(shape_inference::ScalarShape); |
109 | |
110 | REGISTER_OP("AudioSummary" ) |
111 | .Input("tag: string" ) |
112 | .Input("tensor: float" ) |
113 | .Output("summary: string" ) |
114 | .Attr("sample_rate: float" ) |
115 | .Attr("max_outputs: int >= 1 = 3" ) |
116 | .SetShapeFn(shape_inference::ScalarShape) |
117 | .Deprecated(15, "Use AudioSummaryV2." ); |
118 | |
119 | REGISTER_OP("Timestamp" ) |
120 | .Output("ts: float64" ) |
121 | .SetIsStateful() |
122 | .SetShapeFn(shape_inference::ScalarShape); |
123 | |
124 | ALLOW_STATEFUL_OP_FOR_DATASET_FUNCTIONS("Timestamp" ); |
125 | |
126 | } // end namespace tensorflow |
127 | |