Fix false positive tracer leaks in flax library.
We'd like flax itself never to accidentally trigger the jax leak detector when enabled,
but over the years a few longer-lived temporaries were introduced to the flax codebase
that do trigger it. We remedy those cases:
- in flax.core remove a redundant use of nonlocal scopes in pack implementation.
- for nn.jit have the hashing helper datastructure maintain a weakref to module.
- explicitly exempt a context capture used in module summary tabulation routine.
PiperOrigin-RevId: 680626207
