Add overload to generate tensor types with random rank. Also set the max generated dimension to a reasonable value.

LiteRT-PiperOrigin-RevId: 775418960

Author: LukeBoyer

litert/cc/litert_rng.h

@@ -233,15 +233,16 @@ class RandomTensorType {
   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();
+  // Max value of a random dimension.
+  static constexpr auto kMaxDim = 1024;
 
+  // Generate a random tensor type with a pre-determined rank given
+  // in `shape_spec`.
   template <typename Rng>
   Expected<LiteRtRankedTensorType> Generate(
-      Rng& rng,
+      Rng& rng, const std::vector<std::optional<DimSpec>>& shape_spec,
       const ElementTypeSpec& type = {kLiteRtElementTypeInt32,
-                                     kLiteRtElementTypeFloat32},
-      const std::vector<std::optional<DimSpec>>& shape_spec = {}) {
+                                     kLiteRtElementTypeFloat32}) {
     const auto rank = shape_spec.size();
     if (rank > LITERT_TENSOR_MAX_RANK) {
       return Error(kLiteRtStatusErrorInvalidArgument, "Rank too large");
@@ -255,12 +256,24 @@ class RandomTensorType {
     return res;
   }
 
+  // Generate a random tensor type with a random rank.
+  template <typename Rng>
+  Expected<LiteRtRankedTensorType> Generate(
+      Rng& rng, size_t max_rank = LITERT_TENSOR_MAX_RANK,
+      const ElementTypeSpec& type = {kLiteRtElementTypeInt32,
+                                     kLiteRtElementTypeFloat32}) {
+    DimGenerator rank_gen(0, max_rank);
+    std::vector<std::optional<DimSpec>> shape_spec(rank_gen(rng), std::nullopt);
+    return Generate(rng, shape_spec, type);
+  }
+
  private:
   template <typename Rng>
   DimSize GenerateDim(Rng& rng, const std::optional<DimSpec>& dim) {
     if (!dim) {
-      DimGenerator gen;
-      return gen(rng);
+      DimGenerator gen(0, kMaxDim);
+      const auto res = gen(rng);
+      return res;
     } else if (std::holds_alternative<DimSize>(*dim)) {
       auto d = std::get<DimSize>(*dim);
       return d;

litert/cc/litert_rng_test.cc

@@ -156,8 +156,8 @@ TEST_F(LiteRtRngTest, FullySpecifiedRandomTensorType) {
   auto device = TracedDevice();
   RandomTensorType type;
   auto tensor_type = type.Generate(
-      device, {kLiteRtElementTypeFloat32},
-      {RandomTensorType::DimSpec(2u), RandomTensorType::DimSpec(2u)});
+      device, {RandomTensorType::DimSpec(2u), RandomTensorType::DimSpec(2u)},
+      {kLiteRtElementTypeFloat32});
   ASSERT_TRUE(tensor_type);
   EXPECT_EQ(tensor_type->element_type, kLiteRtElementTypeFloat32);
   EXPECT_EQ(tensor_type->layout.dimensions[0], 2);
@@ -168,7 +168,7 @@ TEST_F(LiteRtRngTest, RandomElementType) {
   auto device = TracedDevice();
   RandomTensorType type;
   auto tensor_type = type.Generate(
-      device, {kLiteRtElementTypeFloat32, kLiteRtElementTypeInt32});
+      device, {}, {kLiteRtElementTypeFloat32, kLiteRtElementTypeInt32});
   ASSERT_TRUE(tensor_type);
   EXPECT_TRUE(tensor_type->element_type == kLiteRtElementTypeFloat32 ||
               tensor_type->element_type == kLiteRtElementTypeInt32);
@@ -178,8 +178,8 @@ TEST_F(LiteRtRngTest, RandomTensorShape) {
   auto device = TracedDevice();
   RandomTensorType type;
   auto tensor_type =
-      type.Generate(device, {kLiteRtElementTypeFloat32},
-                    {RandomTensorType::DimRange(1u, 3u), std::nullopt});
+      type.Generate(device, {RandomTensorType::DimRange(1u, 3u), std::nullopt},
+                    {kLiteRtElementTypeFloat32});
   ASSERT_TRUE(tensor_type);
   EXPECT_EQ(tensor_type->element_type, kLiteRtElementTypeFloat32);
   EXPECT_EQ(tensor_type->layout.rank, 2);
@@ -188,7 +188,15 @@ TEST_F(LiteRtRngTest, RandomTensorShape) {
   EXPECT_LE(dim1, 3u);
   const auto dim2 = tensor_type->layout.dimensions[1];
   EXPECT_GE(dim2, 0u);
-  EXPECT_LE(dim2, NumericLimits<uint32_t>::Max());
+  EXPECT_LE(dim2, RandomTensorType::kMaxDim);
+}
+
+TEST_F(LiteRtRngTest, RandomTensorShapeWithRandomRank) {
+  auto device = TracedDevice();
+  RandomTensorType type;
+  auto tensor_type = type.Generate(device, /*max_rank=*/4);
+  ASSERT_TRUE(tensor_type);
+  EXPECT_LE(tensor_type->layout.rank, 4);
 }
 
 }  // namespace