Updates to support jax>0.4.14
Closed this issue · 0 comments
jsbrittain commented
@dfm Firstly, thank you for the write-up which has been very helpful to read through!
I've noticed that the code will no longer run on jax versions > 0.4.14 (currently 0.4.19). Having worked through the issues (which are all fairly minor) I would be happy to submit a PR to update the repo.
Briefly, these consist of:
- deprecated
register_cpu_custom_call_target
function - changes to the api and return values for the jax helper
custom_call
function ShapedArray
has moved tojax.core
for import- I would also like to tag the jax version installed as the colab notebook does not currently run (and the jax api seems to be continuing to evolve)