1/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/core/framework/common_shape_fns.h"
17#include "tensorflow/core/framework/op.h"
18#include "tensorflow/core/framework/rng_alg.h"
19#include "tensorflow/core/framework/shape_inference.h"
20
21namespace tensorflow {
22
23Status StatefulRandomShape(shape_inference::InferenceContext* c) {
24 using shape_inference::ShapeHandle;
25 // Check algorithm shape
26 ShapeHandle unused;
27 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
28 // Set output shape
29 ShapeHandle out;
30 TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &out));
31 c->set_output(0, out);
32 return OkStatus();
33}
34
35#define REGISTER_STATEFUL_OP(name, default_dtype) \
36 REGISTER_OP(name) \
37 .Input("resource: resource") \
38 .Input("algorithm: int64") \
39 .Input("shape: shape_dtype") \
40 .Output("output: dtype") \
41 .Attr("dtype : type = " #default_dtype) \
42 .Attr("shape_dtype : type = DT_INT64") \
43 .SetShapeFn(StatefulRandomShape);
44
45REGISTER_STATEFUL_OP("StatefulUniform", DT_FLOAT);
46REGISTER_STATEFUL_OP("StatefulUniformFullInt", DT_UINT64);
47REGISTER_STATEFUL_OP("StatefulStandardNormalV2", DT_FLOAT);
48REGISTER_STATEFUL_OP("StatefulTruncatedNormal", DT_FLOAT);
49
50REGISTER_OP("StatefulUniformInt")
51 .Input("resource: resource")
52 .Input("algorithm: int64")
53 .Input("shape: shape_dtype")
54 .Input("minval: dtype")
55 .Input("maxval: dtype")
56 .Output("output: dtype")
57 .Attr("dtype : type = DT_INT64")
58 .Attr("shape_dtype : type = DT_INT64")
59 .SetShapeFn([](shape_inference::InferenceContext* c) {
60 using shape_inference::ShapeHandle;
61 // Check inputs
62 ShapeHandle unused;
63 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
64 Status s = c->WithRank(c->input(3), 0, &unused);
65 if (!s.ok()) {
66 return errors::InvalidArgument(
67 "minval must be a scalar; got a tensor of shape ",
68 c->DebugString(c->input(3)));
69 }
70 s = c->WithRank(c->input(4), 0, &unused);
71 if (!s.ok()) {
72 return errors::InvalidArgument(
73 "maxval must be a scalar; got a tensor of shape ",
74 c->DebugString(c->input(4)));
75 }
76 // Set output
77 ShapeHandle out;
78 TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &out));
79 c->set_output(0, out);
80 return OkStatus();
81 });
82
83REGISTER_OP("RngSkip")
84 .Input("resource: resource")
85 .Input("algorithm: int64")
86 .Input("delta: int64")
87 .SetShapeFn([](shape_inference::InferenceContext* c) {
88 shape_inference::ShapeHandle unused;
89 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
90 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
91 return OkStatus();
92 });
93
94REGISTER_OP("RngReadAndSkip")
95 .Input("resource: resource")
96 .Input("alg: int32")
97 .Input("delta: uint64")
98 .Output("value: int64")
99 .SetShapeFn([](shape_inference::InferenceContext* c) {
100 shape_inference::ShapeHandle unused;
101 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
102 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
103 c->set_output(0, c->MakeShape({RNG_MAX_COUNTER_SIZE + RNG_KEY_SIZE}));
104 return OkStatus();
105 });
106
107REGISTER_OP("NonDeterministicInts")
108 .Input("shape: shape_dtype")
109 .SetIsStateful()
110 .Output("output: dtype")
111 .Attr("dtype : type = DT_INT64")
112 .Attr("shape_dtype : type = DT_INT64")
113 .SetShapeFn([](shape_inference::InferenceContext* c) {
114 using shape_inference::ShapeHandle;
115 ShapeHandle out;
116 TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &out));
117 c->set_output(0, out);
118 return OkStatus();
119 });
120
121REGISTER_OP("StatefulRandomBinomial")
122 .Input("resource: resource")
123 .Input("algorithm: int64")
124 .Input("shape: S")
125 .Input("counts: T")
126 .Input("probs: T")
127 .Output("output: dtype")
128 .Attr("S: {int32, int64}")
129 .Attr("T: {half, float, double, int32, int64} = DT_DOUBLE")
130 .Attr("dtype: {half, float, double, int32, int64} = DT_INT64")
131 .SetShapeFn([](shape_inference::InferenceContext* c) {
132 using shape_inference::ShapeHandle;
133
134 ShapeHandle out;
135 TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &out));
136 c->set_output(0, out);
137 return OkStatus();
138 });
139
140// Register the deprecated 'StatefulStandardNormal' op. This op is a short-lived
141// version where the 'resource' variable also contains the algorithm tag.
142// It is deprecated in favor of 'StatefulStandardNormalV2'.
143REGISTER_OP("StatefulStandardNormal")
144 .Deprecated(29, "Use StatefulStandardNormalV2 instead")
145 .Input("resource: resource")
146 .Input("shape: shape_dtype")
147 .Output("output: dtype")
148 .Attr("dtype : type = DT_FLOAT")
149 .Attr("shape_dtype : type = DT_INT64")
150 .SetShapeFn([](shape_inference::InferenceContext* c) {
151 using shape_inference::ShapeHandle;
152 // Set output shape
153 ShapeHandle out;
154 TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &out));
155 c->set_output(0, out);
156 return OkStatus();
157 });
158
159} // namespace tensorflow
160