Use jax.named_scope for name stack rather than named_call.#2349

copybara-service[bot] merged 1 commit intomaingoogle/flax:mainfrom

test_464009184google/flax:test_464009184Copy head branch name to clipboard