From f909ebee9da79da8d35c7bbb1e782e0a0319eefb Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Mon, 10 Feb 2025 13:52:32 -0800
Subject: [PATCH] Fix flaky rand() tests. (#8340)

Summary:

Adjusts the upper bound to prevent rounding to 1.0 when converting to lower-precision types.

Reviewed By: kirklandsign

Differential Revision: D69402222
---
 extension/tensor/tensor_ptr_maker.cpp | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/extension/tensor/tensor_ptr_maker.cpp b/extension/tensor/tensor_ptr_maker.cpp
index ce9de03444c..8e7c908bf43 100644
--- a/extension/tensor/tensor_ptr_maker.cpp
+++ b/extension/tensor/tensor_ptr_maker.cpp
@@ -140,12 +140,22 @@ TensorPtr rand_strided(
     std::vector<executorch::aten::StridesType> strides,
     executorch::aten::ScalarType type,
     executorch::aten::TensorShapeDynamism dynamism) {
+  auto upper_bound = 1.0f;
+  // Adjusts the upper bound to prevent rounding to 1.0 when converting to
+  // lower-precision types.
+  if (type == executorch::aten::ScalarType::Half) {
+    upper_bound -=
+        float(std::numeric_limits<executorch::aten::Half>::epsilon()) / 2;
+  } else if (type == executorch::aten::ScalarType::BFloat16) {
+    upper_bound -=
+        float(std::numeric_limits<executorch::aten::BFloat16>::epsilon()) / 2;
+  }
   return random_strided(
       std::move(sizes),
       std::move(strides),
       type,
       dynamism,
-      std::uniform_real_distribution<float>(0.0f, 1.0f));
+      std::uniform_real_distribution<float>(0.0f, upper_bound));
 }
 
 TensorPtr randn_strided(