Remove uses of deprecated jax.lax.tie_in

google / trax

Trax — Deep Learning with Clear Code and Speed

Apache License 2.0

8.01k stars 813 forks source link

Remove uses of deprecated jax.lax.tie_in #1796

Closed copybara-service[bot] closed 5 months ago

copybara-service[bot] commented 5 months ago

Remove uses of deprecated jax.lax.tie_in

This has been a no-op since jax v0.2.0, and passes the second argument through unchanged. tie_in will be deprecated as of jax v0.4.24; see https://github.com/google/jax/pull/19413