[libneuronxla] Add support for PJRT_Client_CreateUninitializedBuffer

[steeve]

Describe the feature

Sadly it's not possible to use PJRT_Client_CreateUninitializedBuffer.
Still doesn't work on libneuronxla-2.2.8201.0+f46ac1ef

Use Case

This is very useful when creating KV caches.

Proposed Solution

Implement PJRT_Client_CreateUninitializedBuffer using nrt_tensor_allocate, should be fairly trivial ?

Other Information

No response

Acknowledgements

• I may be able to implement this feature request

[github-actions]

Hi @steeve, Thank you for filing the issue! We will take a look and get back to you.

[CarolusSolis]

Hi @steeve,

Thank you for the feature request. We're tracking this internally and will look into adding support for PJRT_Client_CreateUninitializedBuffer. We'll follow up here with updates.