Changelog
Source:NEWS.md
pjrt 0.3.0
Features
- Added
buffer_copy()function to copy buffer between devices. - New
pjrt_register_custom_call()allows external packages to register C/C++ XLA FFI handlers with the PJRT plugin. Registration is deferred until the plugin loads, so handlers can be registered during.onLoad(). -
pjrt_device()now returns cachedPJRTDeviceinstances, so repeated calls for the same device yield objects with stable identity (useful for hashing and caching, e.g. in{anvil}’s JIT).
Bug fixes
- The configure script now uses the
protoccompiler from the same installation as the linked protobuf library, preventing version mismatches when multiple protobuf versions are installed. - Compiling a program for a specific CPU device (e.g.
cpu:1) now targets that device instead of silently falling back tocpu:0. - Fixed device targeting when compiling against a distributed PJRT client, where global device IDs and local hardware ordinals diverge.
pjrt 0.2.0
Asynchronous API
Operations such as host <-> device transfers and program execution were previously only synchronous. Now, they are asynchronous which has considerable performance benefits, especially on GPU. Specifically:
-
pjrt_buffer()andpjrt_execute()return immediately, but the returned buffer is not necessarily ready. To await a transfer or computation of a buffer, useawait(). However, this is handled within PJRT, so this function never has to be called by a user. -
as_array()is still synchronous, but there is now the asynchronous versionas_array_async()but this is rarely needed. If used, it returns aPJRTArrayPromiseobject which can be converted to an Rarray/vectorviavalue(). - To check whether a
PJRTBufferorPJRTArrayPromiseis ready, useis_ready().
Features
- Added
dtypesupport forPJRTBuffers via thetengen::dtypeS3 generic."bool"is now accepted as an alias for"i1"/"pred". - Accept
DataTypeobjects in thedtypeparameter ofpjrt_buffer(). - Support
deviceargument inpjrt_compile().
Bug fixes
- Protect from segfaults in raw to buffer conversion.
- Protect from segfault during device mismatch in
pjrt_execute().