random_ops.cc | Explore in Territory

/* Copyright 2021 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 <algorithm>
#include <cmath>
#include <cstdint>
#include <random>

#include "tensorflow/lite/core/c/builtin_op_data.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tsl/lib/random/philox_random.h"
#include "tsl/lib/random/random_distributions_utils.h"

namespace tflite {
namespace ops {
namespace builtin {
namespace random {

namespace {

Generator;

enum RandomType { … };

struct OpData { … };

// Initialize the OpData based on the seed and seed2 values.
void InitializeOpData(TfLiteNode* node) { … }

// Generates random numbers following a uniform distribution.
// Source: third_party/tensorflow/core/kernels/random_op.cc
void GenerateRandomUniformNumbers(
    Generator& rng, float* buffer, size_t buffer_size) { … }

// Generates random numbers following a standard normal distribution.
// Source: third_party/tensorflow/core/kernels/random_op.cc
void GenerateRandomStandardNormalNumbers(
    Generator& rng, float* buffer, size_t buffer_size) { … }

// Generates random numbers following a multinomial distribution.
// Source: third_party/tensorflow/core/kernels/multinomial_op.cc
template <typename IntType>
void GenerateMultinomialNumbers(Generator& rng, int batch_size,
                                const float* logits, size_t logits_size,
                                IntType* output, size_t num_samples) { … }

}  // namespace

void* Init(TfLiteContext* context, const char* buffer, size_t length) { … }

void Free(TfLiteContext* context, void* buffer) { … }

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { … }

TfLiteStatus PrepareMultinomial(TfLiteContext* context, TfLiteNode* node) { … }

TfLiteStatus EvalRandomType(
    TfLiteContext* context, TfLiteNode* node, RandomType random_type) { … }

template <RandomType rtype>
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { … }

TfLiteStatus EvalMultinomial(TfLiteContext* context, TfLiteNode* node) { … }

}  // namespace random

TfLiteRegistration* Register_RANDOM_UNIFORM() { … }

TfLiteRegistration* Register_RANDOM_STANDARD_NORMAL() { … }

TfLiteRegistration* Register_MULTINOMIAL() { … }

}  // namespace builtin
}  // namespace ops
}  // namespace tflite
chromium/third_party/tflite/src/tensorflow/lite/kernels/random_ops.cc
