Fix a bug where creating a TFRT buffer may access an unset byte_strid…

…es object when the layout is packed. PiperOrigin-RevId: 714992134
shraiysh · Jan 13, 2025 · 135717d · 135717d
1 parent 8d06f36
commit 135717d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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));
       }