[XLA] Fix raw memory Segfault on array queries against default-constructed Buffer shapes

Problem:
When a Buffer shape is instantiated via Shape(BUFFER) without setting its underlying buffer_shape, its internal shape is in InvalidState. If an array property accessor (such as dimensions()) is invoked on this object, the helper array_state_maybe_underneath_buffer() unconditionally dereferenced if_array_state(), causing a fatal Segmentation Fault.

Solution:
Refactored array_state_maybe_underneath_buffer() to call array_state() instead of if_array_state(). This ensures that attempting to read array properties from an uninitialized or broken Buffer shape triggers a graceful, programmer-friendly compiler CHECK failure with clear instructions, rather than executing a raw memory crash.

Testing:
Added unit test BufferShapeArrayAccessorTriggersInformativeCheckFailure in shape_test.cc.
PiperOrigin-RevId: 932772958

Author: tensorflower-gardener

third_party/xla/xla/shape.h

@@ -840,17 +840,31 @@ inline const Shape::ArrayState& Shape::array_state_maybe_underneath_buffer()
     return *array;
   }
   const BufferState* buffer = if_buffer_state();
-  CHECK_NE(buffer, nullptr);
-  return *buffer->buffer_shape->if_array_state();
+  CHECK(buffer)
+      << "Expected an array or buffer shape. Got " << ToString()
+      << "\nThis is a programmer error. Please read the Shape object's "
+         "array properties only when it is an array or buffer shape.";
+  CHECK(buffer->buffer_shape->IsArrayExcludingBuffer())
+      << "Expected a fully initialized Buffer shape containing an "
+         "underlying array, but got an uninitialized/empty Buffer shape: "
+      << ToString();
+  return buffer->buffer_shape->array_state();
 }
 
 inline Shape::ArrayState& Shape::array_state_maybe_underneath_buffer() {
   if (ArrayState* array = if_array_state(); ABSL_PREDICT_TRUE(array)) {
     return *array;
   }
   BufferState* buffer = if_buffer_state();
-  CHECK_NE(buffer, nullptr);
-  return *buffer->buffer_shape->if_array_state();
+  CHECK(buffer)
+      << "Expected an array or buffer shape. Got " << ToString()
+      << "\nThis is a programmer error. Please read the Shape object's "
+         "array properties only when it is an array or buffer shape.";
+  CHECK(buffer->buffer_shape->IsArrayExcludingBuffer())
+      << "Expected a fully initialized Buffer shape containing an "
+         "underlying array, but got an uninitialized/empty Buffer shape: "
+      << ToString();
+  return buffer->buffer_shape->array_state();
 }
 
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE const Shape& Shape::tuple_shapes(

third_party/xla/xla/shape_test.cc

@@ -458,5 +458,35 @@ void BM_TupleShapeEqual(::testing::benchmark::State& state) {
 }
 BENCHMARK(BM_TupleShapeEqual);
 
+TEST(Shape, BufferShapeArrayAccessorTriggersInformativeCheckFailure) {
+  Shape shape(BUFFER);
+  // Verifies that attempting to read array properties from an uninitialized or
+  // invalid Buffer shape triggers an explicit, programmer-friendly CHECK
+  // failure rather than a raw memory fault.
+  EXPECT_DEATH_IF_SUPPORTED(
+      shape.dimensions(),
+      "Expected a fully initialized Buffer shape containing an "
+      "underlying array, but got an uninitialized/empty Buffer shape");
+  EXPECT_DEATH_IF_SUPPORTED(
+      shape.set_dynamic_dimension(0, true),
+      "Expected a fully initialized Buffer shape containing an "
+      "underlying array, but got an uninitialized/empty Buffer shape");
+}
+
+TEST_F(ShapeTest,
+       NonArrayNonBufferShapeAccessorTriggersInformativeCheckFailure) {
+  EXPECT_DEATH_IF_SUPPORTED(tuple_.dimensions(),
+                            "Expected an array or buffer shape.");
+  EXPECT_DEATH_IF_SUPPORTED(token_.dimensions(),
+                            "Expected an array or buffer shape.");
+
+  Shape mutable_tuple = tuple_;
+  EXPECT_DEATH_IF_SUPPORTED(mutable_tuple.set_dynamic_dimension(0, true),
+                            "Expected an array or buffer shape.");
+  Shape mutable_token = token_;
+  EXPECT_DEATH_IF_SUPPORTED(mutable_token.set_dynamic_dimension(0, true),
+                            "Expected an array or buffer shape.");
+}
+
 }  // namespace
 }  // namespace xla