Concretely this is about how BufferFromHostBuffer args may be augmented with metadata for a custom PJRT Plugin without changing the actual arg struct itself here or smuggling metadata through a side channel.
- It seems like PJRT API extensions are the intended path to augment the existing CAPI interface between XLA and a custom PJRT plugin + client instance
- It seems like the path from the c++ PJRTClient -> PjRtCApiClient::BufferFromHostBufferInternalImpl doesn't have a way to attach extensions to the PJRT API call.
So, should I conclude that use of the PJRTClient precludes use of custom PJRT API extensions?
And so, if we wanted to attach some custom extension to the BufferFromHostBuffer args, is it expected that we construct the PJRT API arg struct directly, rather than going through PJRTClient methods as torch_xla does?
An additional note - we are using torch_xla, which instantiates a PJRTClient -> PjRtCApiClient for TPU/custom hardware.
Please correct me if any of these statements are incorrect or there is some other intended mechanism that could be used to augment the current PJRT APIs.
Concretely this is about how BufferFromHostBuffer args may be augmented with metadata for a custom PJRT Plugin without changing the actual arg struct itself here or smuggling metadata through a side channel.
So, should I conclude that use of the PJRTClient precludes use of custom PJRT API extensions?
And so, if we wanted to attach some custom extension to the BufferFromHostBuffer args, is it expected that we construct the PJRT API arg struct directly, rather than going through PJRTClient methods as torch_xla does?
An additional note - we are using torch_xla, which instantiates a PJRTClient -> PjRtCApiClient for TPU/custom hardware.
Please correct me if any of these statements are incorrect or there is some other intended mechanism that could be used to augment the current PJRT APIs.