Add random tensor type generator.

LiteRT-PiperOrigin-RevId: 775336901

Author: LukeBoyer

litert/cc/BUILD

@@ -597,7 +597,11 @@ cc_library(
     ],
     deps = [
         ":litert_detail",
+        ":litert_expected",
         ":litert_numerics",
+        "//litert/c:litert_common",
+        "//litert/c:litert_layout",
+        "//litert/c:litert_model",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
@@ -610,6 +614,7 @@ cc_test(
     deps = [
         ":litert_numerics",
         ":litert_rng",
+        "//litert/c:litert_model_types",
         "//litert/test:rng_fixture",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",

litert/cc/litert_rng.h

@@ -20,11 +20,17 @@
 #include <ostream>
 #include <random>
 #include <type_traits>
+#include <variant>
+#include <vector>
 
 #include "absl/strings/str_cat.h"  // from @com_google_absl
 #include "absl/strings/str_format.h"  // from @com_google_absl
 #include "absl/strings/string_view.h"  // from @com_google_absl
+#include "litert/c/litert_common.h"
+#include "litert/c/litert_layout.h"
+#include "litert/c/litert_model.h"
 #include "litert/cc/litert_detail.h"
+#include "litert/cc/litert_expected.h"
 #include "litert/cc/litert_numerics.h"
 
 // Various utilities and types for random number generation.
@@ -112,7 +118,7 @@ class RandomDevice {
   static constexpr ResultType max() { return Max(); }
 };
 
-// TENSOR DATA GENERATOR ///////////////////////////////////////////////////////
+// PRIMITIVE DATA GENERATORS ///////////////////////////////////////////////////
 
 // Abstract base class for generating data of a certain type from a given rng
 // device, e.g. populating tensors and the like.
@@ -211,6 +217,70 @@ using DefaultRangedGenerator =
 
 using DefaultDevice = RandomDevice<std::mt19937>;
 
+// RANDOM TENSOR TYPES /////////////////////////////////////////////////////////
+
+// This class composes the primitive data generators above to support
+// generating randomized tensor types (and shapes).
+class RandomTensorType {
+ private:
+  using DimSize = uint32_t;
+  using DimGenerator = DefaultRangedGenerator<DimSize>;
+  using ElementTypeInt = uint8_t;
+  using ElementTypeGenerator = DefaultRangedGenerator<ElementTypeInt>;
+
+ public:
+  using DimRange = std::pair<DimSize, DimSize>;
+  using DimSpec = std::variant<DimSize, DimRange>;
+  using ElementTypeSpec = std::vector<LiteRtElementType>;
+
+  static constexpr auto kMinDim = NumericLimits<DimSize>::Lowest();
+  static constexpr auto kMaxDim = NumericLimits<DimSize>::Max();
+
+  template <typename Rng>
+  Expected<LiteRtRankedTensorType> Generate(
+      Rng& rng,
+      const ElementTypeSpec& type = {kLiteRtElementTypeInt32,
+                                     kLiteRtElementTypeFloat32},
+      const std::vector<std::optional<DimSpec>>& shape_spec = {}) {
+    const auto rank = shape_spec.size();
+    if (rank > LITERT_TENSOR_MAX_RANK) {
+      return Error(kLiteRtStatusErrorInvalidArgument, "Rank too large");
+    }
+    LiteRtRankedTensorType res;
+    res.layout.rank = rank;
+    res.element_type = GenerateElementType(rng, type);
+    for (auto i = 0; i < rank; ++i) {
+      res.layout.dimensions[i] = GenerateDim(rng, shape_spec[i]);
+    }
+    return res;
+  }
+
+ private:
+  template <typename Rng>
+  DimSize GenerateDim(Rng& rng, const std::optional<DimSpec>& dim) {
+    if (!dim) {
+      DimGenerator gen;
+      return gen(rng);
+    } else if (std::holds_alternative<DimSize>(*dim)) {
+      auto d = std::get<DimSize>(*dim);
+      return d;
+    } else {
+      auto d = std::get<DimRange>(*dim);
+      DimGenerator gen(d.first, d.second);
+      return gen(rng);
+    }
+  }
+
+  template <typename Rng>
+  LiteRtElementType GenerateElementType(Rng& rng, const ElementTypeSpec& type) {
+    if (type.size() == 1) {
+      return type.front();
+    }
+    ElementTypeGenerator gen(0, type.size() - 1);
+    return type[gen(rng)];
+  }
+};
+
 }  // namespace litert
 
 #endif  // THIRD_PARTY_ODML_LITERT_LITERT_CC_LITERT_RNG_H_

litert/cc/litert_rng_test.cc

@@ -18,13 +18,15 @@
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <optional>
 #include <random>
 #include <type_traits>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/strings/str_format.h"  // from @com_google_absl
 #include "absl/strings/string_view.h"  // from @com_google_absl
+#include "litert/c/litert_model.h"
 #include "litert/cc/litert_numerics.h"
 #include "litert/test/rng_fixture.h"
 
@@ -150,5 +152,44 @@ TEST_F(LiteRtRngTest, TestWithFuzz) {
   }
 }
 
+TEST_F(LiteRtRngTest, FullySpecifiedRandomTensorType) {
+  auto device = TracedDevice();
+  RandomTensorType type;
+  auto tensor_type = type.Generate(
+      device, {kLiteRtElementTypeFloat32},
+      {RandomTensorType::DimSpec(2u), RandomTensorType::DimSpec(2u)});
+  ASSERT_TRUE(tensor_type);
+  EXPECT_EQ(tensor_type->element_type, kLiteRtElementTypeFloat32);
+  EXPECT_EQ(tensor_type->layout.dimensions[0], 2);
+  EXPECT_EQ(tensor_type->layout.dimensions[1], 2);
+}
+
+TEST_F(LiteRtRngTest, RandomElementType) {
+  auto device = TracedDevice();
+  RandomTensorType type;
+  auto tensor_type = type.Generate(
+      device, {kLiteRtElementTypeFloat32, kLiteRtElementTypeInt32});
+  ASSERT_TRUE(tensor_type);
+  EXPECT_TRUE(tensor_type->element_type == kLiteRtElementTypeFloat32 ||
+              tensor_type->element_type == kLiteRtElementTypeInt32);
+}
+
+TEST_F(LiteRtRngTest, RandomTensorShape) {
+  auto device = TracedDevice();
+  RandomTensorType type;
+  auto tensor_type =
+      type.Generate(device, {kLiteRtElementTypeFloat32},
+                    {RandomTensorType::DimRange(1u, 3u), std::nullopt});
+  ASSERT_TRUE(tensor_type);
+  EXPECT_EQ(tensor_type->element_type, kLiteRtElementTypeFloat32);
+  EXPECT_EQ(tensor_type->layout.rank, 2);
+  const auto dim1 = tensor_type->layout.dimensions[0];
+  EXPECT_GE(dim1, 1u);
+  EXPECT_LE(dim1, 3u);
+  const auto dim2 = tensor_type->layout.dimensions[1];
+  EXPECT_GE(dim2, 0u);
+  EXPECT_LE(dim2, NumericLimits<uint32_t>::Max());
+}
+
 }  // namespace
 }  // namespace litert