Enables named call wrapping for labelling profile traces.
When named call wrapping is enabled all JAX ops executed in a Module will be run under
Moduleclass name will show up around the operations belonging to that Module in the Tensorboard profiling UI, simplifying the profiling process.
jax.named_scopeonly works for compiled functions (e.g.: using jax.jit or jax.pmap).