1 /* Copyright 2020 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 <cmath>
17 #include <random>
18
19 #include "tensorflow/lite/c/common.h"
20 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
21 #include "tensorflow/lite/kernels/kernel_util.h"
22
23 namespace tflite {
24 namespace ops {
25 namespace custom {
26 namespace random_standard_normal {
27
28 struct OpData {
29 std::default_random_engine rng;
30 };
31
32 namespace {
33
34 constexpr int kShapeTensor = 0;
35 constexpr int kOutputTensor = 0;
36
37 // Draws a sample from standard normal distribution.
38 template <typename T>
RandomStandardNormalSample(std::default_random_engine & rng,T * output,size_t output_size)39 TfLiteStatus RandomStandardNormalSample(std::default_random_engine& rng,
40 T* output, size_t output_size) {
41 std::normal_distribution<T> dist;
42 std::generate(output, output + output_size, [&]() { return dist(rng); });
43
44 return kTfLiteOk;
45 }
46
47 } // namespace
48
Init(TfLiteContext * context,const char * buffer,size_t length)49 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
50 return new OpData();
51 }
52
Free(TfLiteContext * context,void * buffer)53 void Free(TfLiteContext* context, void* buffer) {
54 delete reinterpret_cast<OpData*>(buffer);
55 }
56
Prepare(TfLiteContext * context,TfLiteNode * node)57 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
58 // The seed/seed2 attributes are not handled in this custom op implementation.
59 TF_LITE_ENSURE(context, NumInputs(node) == 1);
60 TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
61
62 const TfLiteTensor* shape = GetInput(context, node, kShapeTensor);
63 TF_LITE_ENSURE_EQ(context, shape->type, kTfLiteInt32);
64 TF_LITE_ENSURE_EQ(context, NumDimensions(shape), 1);
65 TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
66 if (!IsConstantTensor(shape)) {
67 SetTensorToDynamic(output);
68 return kTfLiteOk;
69 }
70 TfLiteIntArray* output_shape;
71 TF_LITE_ENSURE_OK(context,
72 GetOutputShapeFromInput(context, shape, &output_shape));
73 return context->ResizeTensor(context, output, output_shape);
74 }
75
Eval(TfLiteContext * context,TfLiteNode * node)76 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
77 OpData* params = reinterpret_cast<OpData*>(node->user_data);
78 TF_LITE_ENSURE(context, params != nullptr);
79
80 TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
81 if (IsDynamicTensor(output)) {
82 const TfLiteTensor* shape = GetInput(context, node, kShapeTensor);
83 TfLiteIntArray* output_shape;
84 TF_LITE_ENSURE_OK(context,
85 GetOutputShapeFromInput(context, shape, &output_shape));
86 TF_LITE_ENSURE_OK(context,
87 context->ResizeTensor(context, output, output_shape));
88 }
89 const size_t output_size = NumElements(output);
90 switch (output->type) {
91 case kTfLiteFloat32:
92 RandomStandardNormalSample<float>(
93 params->rng, GetTensorData<float>(output), output_size);
94 break;
95 case kTfLiteFloat64:
96 RandomStandardNormalSample<double>(
97 params->rng, GetTensorData<double>(output), output_size);
98 break;
99 default:
100 TF_LITE_KERNEL_LOG(
101 context, "Unsupported output datatype for RandomStandardNormal: %s",
102 TfLiteTypeGetName(output->type));
103 return kTfLiteError;
104 }
105 return kTfLiteOk;
106 }
107
108 } // namespace random_standard_normal
109
Register_RANDOM_STANDARD_NORMAL()110 TfLiteRegistration* Register_RANDOM_STANDARD_NORMAL() {
111 static TfLiteRegistration r = {
112 random_standard_normal::Init, random_standard_normal::Free,
113 random_standard_normal::Prepare, random_standard_normal::Eval};
114 return &r;
115 }
116
117 } // namespace custom
118 } // namespace ops
119 } // namespace tflite
120