Migrate to new API

[myxik]

Topic says it all, simple fix with just deletes of unnecessary strokes

[xl0]

Hii, @myxik , Thank you for the contribution!

It looks like keeping support for the older versions would be as easy as putting the removed code under a check for the attribute. Could you please do it the way I check for SharedDeviceArray?

Also, I am using nbdev for this project, which means that the library code, CI, and documentation are all generated from Jupiter notebook files (under /nbs/). Could you please

• pip install -e .[dev]
• Bump the version in settings.ini. This looks like a release worthy fix.
• Modify the monkey patch notebook (I use vscode for this). Run the notebook.
• Run nbdev_export (install dev dependencies) to export. I think the notebook itself starts with this command, but you can use it at any time from the cli anyway.
• Run any notebooks that might be affected ( I think just the monkey patch and index ones). Don't forget to restart the kernels to pick up the updated lovely-jax module.
• Run nbdev_prepare - this will run tests and also clean up the notebooks metadata.
• Check that there is no unexpected diff.

[myxik]

Sorry about that, didn't see that the project was on nbdev. Surely, I will try to deliver fixes by today, thanks a lot!