Use jax.closure_convert to make jax_utils.while_loop_bounded with implementation='while_loop' be differentiable through closures. - #2303
Open
copybara-service[bot] wants to merge 1 commit into
Open
