From 135717da248787dd5eaf7a630d02e6053fc44015 Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Mon, 13 Jan 2025 09:11:53 -0800
Subject: [PATCH] Fix a bug where creating a TFRT buffer may access an unset
 byte_strides object when the layout is packed.

PiperOrigin-RevId: 714992134
---
 xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc b/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
index b0b10a6f0e00b..0ec81cb2e3628 100644
--- a/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
+++ b/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
@@ -785,7 +785,9 @@ AbstractTfrtCpuBuffer::BufferFromHostBufferHelper(
         options.elem_size_in_bytes = primitive_util::ByteWidth(type);
         options.dims = dims;
         options.permutation = permutation;
-        options.input_layout = TransposePlan::Striding{*byte_strides};
+        if (byte_strides) {
+          options.input_layout = TransposePlan::Striding{*byte_strides};
+        }
         absl::MutexLock lock(transpose_mu);
         TF_ASSIGN_OR_RETURN(transpose, transpose_cache->GetOrCreate(options));
       }