Check failed: on_device_shape().has_layout()

[SobhanMP]

Hi, I'm not sure if this a bug or not but I've been running into

F0714 01:41:52.913772 634916 pjrt_client.h:980] Check failed: on_device_shape().has_layout()
*** Check failure stack trace: ***
@ 0x7f54dfd15754 absl::lts_20230802::log_internal::LogMessage::SendToLog()
@ 0x7f54dfd155c4 absl::lts_20230802::log_internal::LogMessage::Flush()
@ 0x7f54dfd15af9 absl::lts_20230802::log_internal::LogMessageFatal::~LogMessageFatal()
@ 0x7f54d7a6d9b2 xla::PjRtBuffer::layout()
@ 0x7f54d79ccb80 pjrt::PJRT_Layouts_PJRT_Buffer_MemoryLayout()
@ 0x7f55cfcd2442 pjrt::GetMemoryLayout()
@ 0x7f55cfcc001d xla::PjRtCApiBuffer::layout()
@ 0x7f55d7ccf3cb xla::ifrt::PjRtLoadedExecutable::Execute()
@ 0x7f55cfc29409 jax::(anonymous namespace)::PjitFunction::Call()
@ 0x7f55cfc2580f PjitFunction_tp_vectorcall
@ 0x5596f35ab147 PyObject_Vectorcall
Fatal Python error: Aborted

but I'm having a hard time creating a repro, any tips for finding a root cause?

I'm using :
jax: 0.6.2
jaxlib: 0.6.2
device info: NVIDIA GeForce RTX 4090-1, 1 local devices"