Use jax.named_scope for name stack rather than named_call.
Removes old function-transform-based implementation of named call in favor of
new (post 0.3.14 JAX) jax.named_scope context manager mechanism. Refactors some
tests.
PiperOrigin-RevId: 464144791
