[XLA:GPU] Fix CHECK-fail for tuples of int4.

PiperOrigin-RevId: 665208121
Intel-tensorflow · Aug 20, 2024 · 14f860e · 14f860e
1 parent 2ca3e97
commit 14f860e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc
@@ -252,7 +252,7 @@ Literal::Literal(const Shape& shape)
 void Literal::SetShape(const Shape& shape) {
   Shape shape_storage;
   const Shape* shape_ptr = &shape;
-  if (LayoutUtil::HasCustomElementSizeInBits(shape)) {
+  if (shape.IsArray() && LayoutUtil::HasCustomElementSizeInBits(shape)) {
     shape_storage = shape;
     shape_storage.mutable_layout()->set_element_size_in_bits(0);
     shape_ptr = &shape_storage;

diff --git a/third_party/xla/xla/literal_test.cc b/third_party/xla/xla/literal_test.cc
@@ -2583,6 +2583,14 @@ TEST_F(LiteralUtilTest, CreateFromShapeWithUnknownLeafArrays) {
   EXPECT_FALSE(c1.IsKnown());
 }
 
+TEST_F(LiteralUtilTest, CreateFromShapeWithUnknownLeafArraysS4Tuple) {
+  auto inner_shape = ShapeUtil::MakeShape(S4, {4, 4});
+  inner_shape.mutable_layout()->set_element_size_in_bits(4);
+  Literal c1 = Literal::CreateFromShapeWithUnknownLeafArrays(
+      ShapeUtil::MakeTupleShape({inner_shape}));
+  EXPECT_FALSE(c1.IsKnown());
+}
+
 TEST_F(LiteralUtilTest, CreatePartiallyKnownTuple) {
   Literal c1 = Literal::CreateFromShapeWithUnknownLeafArrays(
       ShapeUtil::MakeShape(F32, {4, 4}));