Remove uses of deprecated jax.lax.tie_in#1795

Open

copybara-service[bot] wants to merge 1 commit intomasterfrom test_602417629

+316-316

Commits on Jan 29, 2024

Remove uses of deprecated jax.lax.tie_in

Jake VanderPlas
authored and
copybara-github
committed