jax

f168a156 - [GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.

[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call. May fix flaky failures in CI. Make stream argument to Pool::Borrow() mandatory to minimize chance of forgetting it. PiperOrigin-RevId: 530425766