aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/ops/stateless_random_ops.cc
blob: b222b5b2416dfac09e1dd1abd15862e317b064c7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"

namespace tensorflow {

using shape_inference::DimensionHandle;
using shape_inference::ShapeHandle;

static Status StatelessShape(shape_inference::InferenceContext* context) {
  // Check seed shape
  ShapeHandle seed;
  TF_RETURN_IF_ERROR(context->WithRank(context->input(1), 1, &seed));
  DimensionHandle unused;
  TF_RETURN_IF_ERROR(context->WithValue(context->Dim(seed, 0), 2, &unused));

  // Set output shape
  shape_inference::ShapeHandle out;
  TF_RETURN_IF_ERROR(context->MakeShapeFromShapeTensor(0, &out));
  context->set_output(0, out);
  return Status::OK();
}

#define REGISTER_STATELESS_OP(name)                  \
  REGISTER_OP(name)                                  \
      .Input("shape: T")                             \
      .Input("seed: int64")                          \
      .Output("output: dtype")                       \
      .Attr("dtype: {half,float,double} = DT_FLOAT") \
      .Attr("T: {int32, int64} = DT_INT32")          \
      .SetShapeFn(StatelessShape)

// This op is exposed through contrib/stateless only.  The interface may change.
REGISTER_STATELESS_OP("StatelessRandomUniform").Doc(R"doc(
Outputs deterministic pseudorandom random values from a uniform distribution.

The generated values follow a uniform distribution in the range `[0, 1)`. The
lower bound 0 is included in the range, while the upper bound 1 is excluded.

The outputs are a deterministic function of `shape` and `seed`.

shape: The shape of the output tensor.
dtype: The type of the output.
seed: 2 seeds (shape [2]).
output: Random values with specified shape.
)doc");

// This op is exposed through contrib/stateless only.  The interface may change.
REGISTER_STATELESS_OP("StatelessRandomNormal").Doc(R"doc(
Outputs deterministic pseudorandom values from a normal distribution.

The generated values will have mean 0 and standard deviation 1.

The outputs are a deterministic function of `shape` and `seed`.

shape: The shape of the output tensor.
dtype: The type of the output.
seed: 2 seeds (shape [2]).
output: Random values with specified shape.
)doc");

// This op is exposed through contrib/stateless only.  The interface may change.
REGISTER_STATELESS_OP("StatelessTruncatedNormal").Doc(R"doc(
Outputs deterministic pseudorandom values from a truncated normal distribution.

The generated values follow a normal distribution with mean 0 and standard
deviation 1, except that values whose magnitude is more than 2 standard
deviations from the mean are dropped and re-picked.

The outputs are a deterministic function of `shape` and `seed`.

shape: The shape of the output tensor.
dtype: The type of the output.
seed: 2 seeds (shape [2]).
output: Random values with specified shape.
)doc");

#undef REGISTER_STATELESS_OP

}  // namespace tensorflow