Snap for 7236705 from ef3bdfa83d72f19474b018177c34e27f0beb4981 to sc-releaseandroid-vts-12.0_r1android-security-12.0.0_r59android-security-12.0.0_r58android-security-12.0.0_r57android-security-12.0.0_r56android-security-12.0.0_r55android-security-12.0.0_r54android-security-12.0.0_r53android-security-12.0.0_r52android-security-12.0.0_r51android-security-12.0.0_r50android-security-12.0.0_r49android-security-12.0.0_r48android-security-12.0.0_r47android-security-12.0.0_r46android-security-12.0.0_r45android-security-12.0.0_r44android-security-12.0.0_r43android-security-12.0.0_r42android-security-12.0.0_r41android-security-12.0.0_r40android-security-12.0.0_r39android-security-12.0.0_r38android-security-12.0.0_r37android-security-12.0.0_r36android-security-12.0.0_r35android-security-12.0.0_r34android-platform-12.0.0_r9android-platform-12.0.0_r8android-platform-12.0.0_r7android-platform-12.0.0_r6android-platform-12.0.0_r5android-platform-12.0.0_r4android-platform-12.0.0_r31android-platform-12.0.0_r30android-platform-12.0.0_r3android-platform-12.0.0_r29android-platform-12.0.0_r28android-platform-12.0.0_r27android-platform-12.0.0_r26android-platform-12.0.0_r25android-platform-12.0.0_r24android-platform-12.0.0_r23android-platform-12.0.0_r22android-platform-12.0.0_r21android-platform-12.0.0_r20android-platform-12.0.0_r2android-platform-12.0.0_r19android-platform-12.0.0_r18android-platform-12.0.0_r17android-platform-12.0.0_r16android-platform-12.0.0_r15android-platform-12.0.0_r14android-platform-12.0.0_r13android-platform-12.0.0_r12android-platform-12.0.0_r11android-platform-12.0.0_r10android-platform-12.0.0_r1android-cts-12.0_r1android-12.0.0_r9android-12.0.0_r8android-12.0.0_r34android-12.0.0_r33android-12.0.0_r31android-12.0.0_r30android-12.0.0_r3android-12.0.0_r25android-12.0.0_r2android-12.0.0_r11android-12.0.0_r10android-12.0.0_r1android12-security-releaseandroid12-s5-releaseandroid12-s4-releaseandroid12-s3-releaseandroid12-s2-releaseandroid12-s1-releaseandroid12-releaseandroid12-platform-release

Change-Id: I468f3203926785d5733cbb60110da62856ab827e

6 files changed, 377 insertions, 0 deletions

diff --git a/tensorflow/lite/kernels/Android.bp b/tensorflow/lite/kernels/Android.bp
index 325daf3185c..94f62282946 100644
--- a/tensorflow/lite/kernels/Android.bp
+++ b/tensorflow/lite/kernels/Android.bp
@@ -127,6 +127,7 @@ cc_library_static {
         "pooling.cc",
         "pow.cc",
         "quantize.cc",
+        "random_uniform.cc",
         "range.cc",
         "rank.cc",
         "reduce.cc",
diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD
index 3db72c79d40..fc9d5c59e87 100644
--- a/tensorflow/lite/kernels/BUILD
+++ b/tensorflow/lite/kernels/BUILD
@@ -734,6 +734,7 @@ cc_library(
     srcs = [
         "multinomial.cc",
         "random_standard_normal.cc",
+        "random_uniform.cc",
     ],
     hdrs = ["custom_ops_register.h"],
     copts = tflite_copts(),
@@ -1442,6 +1443,20 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "random_uniform_test",
+    size = "small",
+    srcs = ["random_uniform_test.cc"],
+    deps = [
+        ":custom_ops",
+        ":test_main",
+        ":test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/testing:util",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 cc_library(
     name = "reshape_test_common",
     testonly = 1,
diff --git a/tensorflow/lite/kernels/custom_ops_register.h b/tensorflow/lite/kernels/custom_ops_register.h
index bbfc6a823c0..3b4ca3d8240 100644
--- a/tensorflow/lite/kernels/custom_ops_register.h
+++ b/tensorflow/lite/kernels/custom_ops_register.h
@@ -27,6 +27,8 @@ TfLiteRegistration* Register_HASHTABLE_IMPORT();
 TfLiteRegistration* Register_HASHTABLE_SIZE();
 TfLiteRegistration* Register_MULTINOMIAL();
 TfLiteRegistration* Register_RANDOM_STANDARD_NORMAL();
+TfLiteRegistration* Register_RANDOM_UNIFORM();
+TfLiteRegistration* Register_RANDOM_UNIFORM_INT();
 
 }  // namespace custom
 }  // namespace ops
diff --git a/tensorflow/lite/kernels/random_uniform.cc b/tensorflow/lite/kernels/random_uniform.cc
new file mode 100644
index 00000000000..812953c7d6d
--- /dev/null
+++ b/tensorflow/lite/kernels/random_uniform.cc
@@ -0,0 +1,181 @@
+/* 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 <limits>
+#include <random>
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace random_uniform {
+
+struct OpData {
+  // This implementation uses a random generator from the standard C++ library
+  // on the platform where TFLite is build. This is different from the TF
+  // version of the kernel that uses custom implementations of random
+  // generator, different for different hardware.
+  std::default_random_engine rng;
+};
+
+namespace {
+
+template <typename T, typename dist_type>
+void RandomUniformSample(std::default_random_engine& rng, T* buffer,
+                         size_t buffer_size, T min_value, T max_value) {
+  dist_type dist(min_value, max_value);
+  std::generate(buffer, buffer + buffer_size, [&]() { return dist(rng); });
+}
+
+TfLiteIntArray* CreateDimensionsFromTensor(const TfLiteTensor* tensor) {
+  const int output_dims = tflite::SizeOfDimension(tensor, 0);
+  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(output_dims);
+  for (int i = 0; i < output_dims; i++) {
+    output_shape->data[i] = tensor->data.i32[i];
+  }
+  return output_shape;
+}
+}  // namespace
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  return new OpData();
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  // TODO(b/169611265): Handle optional seed input.
+  TF_LITE_ENSURE(context, tflite::NumInputs(node) >= 1);
+  TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
+
+  // Input is a shape tensor.
+  const TfLiteTensor* input = tflite::GetInput(context, node, 0);
+  TF_LITE_ENSURE_EQ(context, tflite::NumDimensions(input), 1);
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+  if (!IsConstantTensor(input)) {
+    SetTensorToDynamic(output);
+    return kTfLiteOk;
+  }
+  return context->ResizeTensor(context, output,
+                               CreateDimensionsFromTensor(input));
+}
+
+TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node) {
+  OpData* params = reinterpret_cast<OpData*>(node->user_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+  if (IsDynamicTensor(output)) {
+    const TfLiteTensor* input = tflite::GetInput(context, node, 0);
+    TF_LITE_ENSURE_OK(context,
+                      context->ResizeTensor(context, output,
+                                            CreateDimensionsFromTensor(input)));
+  }
+  const size_t output_size = tflite::NumElements(output);
+  switch (output->type) {
+    case kTfLiteFloat32:
+      RandomUniformSample<float, std::uniform_real_distribution<float>>(
+          params->rng, GetTensorData<float>(output), output_size, 0.f, 1.f);
+      break;
+    case kTfLiteFloat64:
+      RandomUniformSample<double, std::uniform_real_distribution<double>>(
+          params->rng, GetTensorData<double>(output), output_size, 0.f, 1.f);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Unsupported output datatype for RandomUniform: %s",
+                         TfLiteTypeGetName(output->type));
+      return kTfLiteError;
+  }
+
+  return kTfLiteOk;
+}
+
+int64_t IntValueFromTensor(const TfLiteTensor* tensor) {
+  switch (tensor->type) {
+    case kTfLiteInt8:
+      return *GetTensorData<int8_t>(tensor);
+    case kTfLiteInt32:
+      return *GetTensorData<int32_t>(tensor);
+    case kTfLiteInt64:
+      return *GetTensorData<int64_t>(tensor);
+    default:
+      return -1;
+  }
+}
+
+TfLiteStatus EvalInt(TfLiteContext* context, TfLiteNode* node) {
+  OpData* params = reinterpret_cast<OpData*>(node->user_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  TF_LITE_ENSURE(context, tflite::NumInputs(node) >= 3);
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+  if (IsDynamicTensor(output)) {
+    const TfLiteTensor* input = tflite::GetInput(context, node, 0);
+    TF_LITE_ENSURE_OK(context,
+                      context->ResizeTensor(context, output,
+                                            CreateDimensionsFromTensor(input)));
+  }
+  int64_t min_value = IntValueFromTensor(tflite::GetInput(context, node, 1));
+  int64_t max_value = IntValueFromTensor(tflite::GetInput(context, node, 2));
+  TF_LITE_ENSURE(context, min_value < max_value);
+  size_t output_size = tflite::NumElements(output);
+  switch (output->type) {
+    case kTfLiteInt8:
+      RandomUniformSample<int8_t, std::uniform_int_distribution<int32_t>>(
+          params->rng, GetTensorData<int8_t>(output), output_size, min_value,
+          max_value);
+      break;
+    case kTfLiteInt32:
+      RandomUniformSample<int32_t, std::uniform_int_distribution<int32_t>>(
+          params->rng, GetTensorData<int32_t>(output), output_size, min_value,
+          max_value);
+      break;
+    case kTfLiteInt64:
+      RandomUniformSample<int64_t, std::uniform_int_distribution<int64_t>>(
+          params->rng, GetTensorData<int64_t>(output), output_size, min_value,
+          max_value);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Unsupported output datatype for RandomUniformInt: %s",
+                         TfLiteTypeGetName(output->type));
+      return kTfLiteError;
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace random_uniform
+
+TfLiteRegistration* Register_RANDOM_UNIFORM() {
+  static TfLiteRegistration r = {random_uniform::Init, random_uniform::Free,
+                                 random_uniform::Prepare,
+                                 random_uniform::EvalFloat};
+  return &r;
+}
+
+TfLiteRegistration* Register_RANDOM_UNIFORM_INT() {
+  static TfLiteRegistration r = {random_uniform::Init, random_uniform::Free,
+                                 random_uniform::Prepare,
+                                 random_uniform::EvalInt};
+  return &r;
+}
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
diff --git a/tensorflow/lite/kernels/random_uniform_test.cc b/tensorflow/lite/kernels/random_uniform_test.cc
new file mode 100644
index 00000000000..e5c64ed38a7
--- /dev/null
+++ b/tensorflow/lite/kernels/random_uniform_test.cc
@@ -0,0 +1,177 @@
+/* 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 <cstdint>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/custom_ops_register.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+namespace {
+template <typename T>
+tflite::TensorType GetTTEnum();
+
+template <>
+tflite::TensorType GetTTEnum<float>() {
+  return tflite::TensorType_FLOAT32;
+}
+
+template <>
+tflite::TensorType GetTTEnum<double>() {
+  return tflite::TensorType_FLOAT64;
+}
+
+template <>
+tflite::TensorType GetTTEnum<int8_t>() {
+  return tflite::TensorType_INT8;
+}
+
+template <>
+tflite::TensorType GetTTEnum<int32_t>() {
+  return tflite::TensorType_INT32;
+}
+
+template <>
+tflite::TensorType GetTTEnum<int64_t>() {
+  return tflite::TensorType_INT64;
+}
+
+class RandomUniformOpModel : public tflite::SingleOpModel {
+ public:
+  RandomUniformOpModel(const std::initializer_list<int>& input,
+                       tflite::TensorData output, bool dynamic_input) {
+    if (dynamic_input) {
+      input_ = AddInput({tflite::TensorType_INT32, {3}});
+    } else {
+      input_ = AddConstInput(tflite::TensorType_INT32, input,
+                             {static_cast<int>(input.size())});
+    }
+    output_ = AddOutput(output);
+    SetCustomOp("RandomUniform", {}, ops::custom::Register_RANDOM_UNIFORM);
+    BuildInterpreter({GetShape(input_)});
+    if (dynamic_input) {
+      PopulateTensor<int32_t>(input_, std::vector<int32_t>(input));
+    }
+  }
+
+  int input_;
+  int output_;
+
+  int input() { return input_; }
+  int output() { return output_; }
+
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+};
+
+class RandomUniformIntOpModel : public tflite::SingleOpModel {
+ public:
+  RandomUniformIntOpModel(const std::initializer_list<int>& input,
+                          tflite::TensorData output, int min_val, int max_val) {
+    input_ = AddConstInput(tflite::TensorType_INT32, input,
+                           {static_cast<int>(input.size())});
+    input_minval_ = AddConstInput(tflite::TensorType_INT32, {min_val}, {1});
+    input_maxval_ = AddConstInput(tflite::TensorType_INT32, {max_val}, {1});
+    output_ = AddOutput(output);
+    SetCustomOp("RandomUniformInt", {},
+                ops::custom::Register_RANDOM_UNIFORM_INT);
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  int input_;
+  int input_minval_;
+  int input_maxval_;
+
+  int output_;
+
+  int input() { return input_; }
+  int output() { return output_; }
+
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+};
+
+}  // namespace
+}  // namespace tflite
+
+template <typename FloatType>
+class RandomUniformTest : public ::testing::Test {
+ public:
+  using Float = FloatType;
+};
+
+using TestTypes = ::testing::Types<float, double>;
+
+TYPED_TEST_SUITE(RandomUniformTest, TestTypes);
+
+TYPED_TEST(RandomUniformTest, TestOutput) {
+  using Float = typename TestFixture::Float;
+  for (const auto dynamic : {true, false}) {
+    tflite::RandomUniformOpModel m({1000, 50, 5},
+                                   {tflite::GetTTEnum<Float>(), {}}, dynamic);
+    m.Invoke();
+    auto output = m.GetOutput<Float>();
+    EXPECT_EQ(output.size(), 1000 * 50 * 5);
+
+    double sum = 0;
+    for (const auto r : output) {
+      sum += r;
+    }
+    double avg = sum / output.size();
+    ASSERT_LT(std::abs(avg - 0.5), 0.05);  // Average should approximately 0.5
+
+    double sum_squared = 0;
+    for (const auto r : output) {
+      sum_squared += std::pow(r - avg, 2);
+    }
+    double var = sum_squared / output.size();
+    EXPECT_LT(std::abs(1. / 12 - var),
+              0.05);  // Variance should be approximately 1./12
+  }
+}
+
+template <typename IntType>
+class RandomUniformIntTest : public ::testing::Test {
+ public:
+  using Int = IntType;
+};
+
+using TestTypesInt = ::testing::Types<int8_t, int32_t, int64_t>;
+
+TYPED_TEST_SUITE(RandomUniformIntTest, TestTypesInt);
+
+TYPED_TEST(RandomUniformIntTest, TestOutput) {
+  using Int = typename TestFixture::Int;
+  tflite::RandomUniformIntOpModel m({1000, 50, 5},
+                                    {tflite::GetTTEnum<Int>(), {}}, 0, 5);
+  m.Invoke();
+  auto output = m.GetOutput<Int>();
+  EXPECT_EQ(output.size(), 1000 * 50 * 5);
+
+  int counters[] = {0, 0, 0, 0, 0, 0};
+  for (const auto r : output) {
+    ASSERT_GE(r, 0);
+    ASSERT_LE(r, 5);
+    ++counters[r];
+  }
+  // Check that all numbers are meet with near the same frequency.
+  for (int i = 1; i < 6; ++i) {
+    EXPECT_LT(std::abs(counters[i] - counters[0]), 1000);
+  }
+}
diff --git a/tensorflow/lite/tflite_static.bp b/tensorflow/lite/tflite_static.bp
index 6f5c609eeea..7d8943d9ea9 100644
--- a/tensorflow/lite/tflite_static.bp
+++ b/tensorflow/lite/tflite_static.bp
@@ -109,6 +109,7 @@ cc_library_static {
         "kernels/pooling.cc",
         "kernels/pow.cc",
         "kernels/quantize.cc",
+        "kernels/random_uniform.cc",
         "kernels/range.cc",
         "kernels/rank.cc",
         "kernels/reduce.cc",