If the dynamic_dimensions parameter is empty in the Shape ctor, assume all dimensions are static.

Some callers call the `Shape(element_type, dimensions, dynamic_dimensions)` ctor with a non-empty `dimensions` and an empty `dynamic_dimensions`. This breaks the shape object's invariant that the two should have the same size.

We have two options for fixing this:

1. Force the caller to always provide a `dynamic_dimensions` whose size matches that of `dimensions`.
2. Provide a sensible default behavior when `dynamic_dimensions` is empty.

I chose #2 as:

1. #1 is more risky as it may cause the compiler to crash in production (e.g. if we don't have adequate test coverage).
2. It's very common for an array to have only static dimensions. Therefore it's good to optimize the user experience for this common case.

PiperOrigin-RevId: 739197635

Author: Google-ML-Automation

xla/shape.cc

@@ -59,7 +59,10 @@ Shape::Shape(const PrimitiveType element_type,
                           dynamic_dimensions.end()) {
   CHECK(primitive_util::IsArrayType(element_type_))
       << "Invalid element type for array shape: " << element_type_;
-  if (!dynamic_dimensions.empty()) {
+  if (dynamic_dimensions_.empty()) {
+    // Assume all dimensions are static.
+    dynamic_dimensions_.resize(dimensions_.size(), false);
+  } else {
     CHECK_EQ(dimensions_.size(), dynamic_dimensions_.size())
         << "If dynamic_dimensions is provided, it must have the same size as "
            "dimensions.";

xla/shape.h

@@ -67,7 +67,7 @@ class Shape {
   // Precondition:
   //  - `element_type` must be a valid array type.
   //  - `dynamic_dimensions` must be either empty or have the same size as
-  //    `dimensions`.
+  //    `dimensions`. If it's empty, all dimensions are static.
   Shape(PrimitiveType element_type, absl::Span<const int64_t> dimensions,
         absl::Span<const bool> dynamic_dimensions);
 

xla/shape_test.cc

@@ -46,6 +46,16 @@ class ShapeTest : public ::testing::Test {
       ShapeUtil::MakeShape(F32, {Shape::kUnboundedSize, 784}, {true, false});
 };
 
+// Tests that if the dynamic_dimensions parameter empty in the Shape
+// constructor, it's treated as all dimensions are static.
+TEST(Shape, ArrayCtorTreatsEmptyDynamicDimensionsAsAllStatic) {
+  const Shape shape(F32, {1, 2, 3}, {});
+  EXPECT_TRUE(shape.is_static());
+  EXPECT_TRUE(shape.is_static_dimension(0));
+  EXPECT_TRUE(shape.is_static_dimension(1));
+  EXPECT_TRUE(shape.is_static_dimension(2));
+}
+
 TEST_F(ShapeTest, ShapeToFromProto) {
   for (const Shape& shape :
        {opaque_, token_, scalar_, matrix_, matrix2_, tuple_, nested_tuple_,