search

google / trax

Trax — Deep Learning with Clear Code and Speed
Apache License 2.0
8.01k stars 813 forks source link

Replace references to deprecated device_buffer attributes #1793

Open copybara-service[bot] opened 7 months ago

copybara-service[bot] commented 7 months ago

Replace references to deprecated device_buffer attributes

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.