1 /* Copyright 2019 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 #ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_RANDOM_OPS_UTIL_H_
17 #define TENSORFLOW_COMPILER_TF2XLA_KERNELS_RANDOM_OPS_UTIL_H_
18 
19 #include <cmath>
20 
21 #include "tensorflow/compiler/xla/client/xla_builder.h"
22 #include "tensorflow/core/framework/tensor.h"
23 
24 namespace tensorflow {
25 // Returns a tensor containing 'shape' random values uniformly distributed in
26 // the range [minval, maxval). The raw random bits are generated by the given
27 // `bit_generator` and converted to the requested data type and range. This
28 // routine requires 2 32-bit integer seeds and currently only supports 'shape's
29 // of type F32, S32 and S64.
30 xla::XlaOp StatelessRngUniform(absl::string_view device_type_string,
31                                xla::XlaOp seeds, const xla::Shape& shape,
32                                xla::XlaOp minval, xla::XlaOp maxval);
33 
34 // Converts to bfloat16 if `dtype` equals DT_BFLOAT16, no-op otherwise.
35 // It masks the last 16 bit. With normal rounding, values near "maxval" would be
36 // converted to "maxval" which is out of range ["minval", "maxval"). In
37 // addition, the distribution near the limit is not uniform.
38 xla::XlaOp MaybeConvertF32ToBF16(xla::XlaOp input, DataType dtype);
39 
40 // Combines two signed 32-bit seeds into a single unsigned 64 bit seed.
41 xla::XlaOp GetU64FromS32Seeds(xla::XlaOp seed0, xla::XlaOp seed1);
42 
43 }  // namespace tensorflow
44 
45 #endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_RANDOM_OPS_UTIL_H_
46