Update named_call so that it still supports Python control flow after JAX's omnistaging change (https://github.com/google/jax/pull/3370).
Update named_call so that it still supports Python control flow after JAX's
omnistaging change (https://github.com/google/jax/pull/3370).

PiperOrigin-RevId: 326651898
Change-Id: I177f24632c3e1fb802fdf218c4d117ca8da5308f

deepmind/dm-haiku,deepmind/dm-haiku