Add a check that if any single dimension of an xla shape is 0 the dat…

…a size is 0.

PiperOrigin-RevId: 636352383

third_party/xla/xla/shape_util_test.cc

@@ -1202,6 +1202,20 @@ TEST(ShapeUtilTest, Int4ShapeSize) {
   EXPECT_EQ(ShapeUtil::ArraySize(int4_shape2), 9216 * 6144 / 2);
 }
 
+TEST(XlaShapeUtilTest, ZeroSize) {
+  // Verify that if any one dimension is 0 we have a zero byte buffer.
+  std::vector<std::vector<int64_t>> test_cases = {
+      {0, 64, 128}, {128, 0, 64}, {64, 128, 0},
+      {0, 63, 127}, {127, 0, 63}, {63, 127, 0},
+  };
+  for (const auto& dimensions : test_cases) {
+    xla::Shape int4_shape = xla::ShapeUtil::MakeShape(xla::S4, dimensions);
+    int4_shape.mutable_layout()->set_element_size_in_bits(4);
+    EXPECT_EQ(xla::ShapeUtil::ArrayDataSize(int4_shape), 0);
+    EXPECT_EQ(xla::ShapeUtil::ArraySize(int4_shape), 0);
+  }
+}
+
 TEST(ShapeUtilTest, DecomposeBitcastToReshape) {
   const Shape kInputShape =
       ShapeUtil::MakeShapeWithDenseLayout(F32, {1, 16, 17, 3}, {3, 2, 1, 0});