BufferFromHostBuffer for pre-allocated GPU buffers

[drewjenks01]

I have built XLA from source using the version Jax 0.4.25 (4ccfe33) uses and am calling the BufferFromHostBuffer function here:

xla/xla/pjrt/pjrt_stream_executor_client.cc

Line 834 in 4ccfe33

bool is_cpu_platform =

.

I see that there is a fast-path for when you are on CPU and passing a CPU buffer, but there is no option for the same thing with GPU buffers. I want to be able to pass in a pre-allocated GPU buffer and not have XLA make a copy, but rather just point to this buffer like it does with CPU buffers. I'm wondering why this isn't a feature within this function or if this has changed since the version I'm using? I'm considering building in the path myself but wondering if there will be any negative side effects. Thanks!