Replace references to deprecated device_buffer attributes

Open copybara-service[bot] opened this issue 2 years ago • 1 comments

jax.Array.device_buffer and jax.Array.device_buffers will be deprecated as of jax version 0.4.22; see https://github.com/google/jax/pull/18844.

Dec 06 '23 23:12 copybara-service[bot]

@niketkumar could you please consider merging this. Indeed those methods were definitively removed from jax in the latest release (0.4.27).

May 07 '24 22:05 GaetanLepage