Layout API propagation across PJRT C-APIs

[Devesh Ratho]

Hello everyone,

I'm trying to understand the propagation of Layout information across PJRT C-APIs.

My current understanding is that PJRT_Compile has access to layout information via PJRT_Program (via the input mlir module).

For the stream executor client, layout information is associated with buffers via CreateBuffersForAsyncHostToDevice, but this API is not supported for the C-API client. How is device layout information propagated to other APIs, such as BufferFromHostBuffer, that have access to layout information?

Thanks,

Devesh