Add layer-by-layer tensor distribution debugging mode in MaxText.

[copybara-service]

Add layer-by-layer tensor distribution debugging mode in MaxText.

This change introduces:

1. Core Tensor Telemetry (debug_tensor_utils.py): Pure JAX/NumPy utility logging tensor distributions (mean, std, min, max, L2 norm, percentiles [1%, 5%, 25%, 50%, 75%, 95%, 99%], NaN count, Inf count, and continuous MoE expert routing mass vectors) across both Forward (FWD) and Backward (BWD) passes.
2. Interceptor Telemetry Infrastructure (debug_tensor_interceptors.py): Modular Linen interceptor (linen_interceptor_fn) and NNX module wrapper (wrap_nnx_module_for_debug) paired with thread-local context management (debug_telemetry_scope). Automatically constructs hierarchical scope paths (e.g. decoder/layers_0/self_attention/query, decoder/layers_0/MoeBlock_0/router_weights) with zero changes to individual model and layer classes.
3. Centralized NNX Hook in ToLinen (layers/nnx_wrappers.py): Integrates with Flax Linen and NNX hybrid architectures.
4. FWD and BWD Pass Autodiff: Uses @jax.custom_vjp with static trace-time short-circuiting to automatically log forward activations and reverse-mode gradient cotangents with zero overhead when disabled.

Tested:

• Unit tests in //third_party/py/maxtext/tests/unit:debug_tensor_utils_test.
• Integration tests in //third_party/py/maxtext/tests/unit:debug_tensor_integration_test.
• End-to-end synthetic Mixtral MoE training run (trainers:train with model_name=mixtral-8x7b, num_experts=4, num_experts_per_tok=2, debug_tensor_distribution=True):

  [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/expert_inputs shape=(12, 128, 64) dtype=bfloat16 mean=-8.294509e-03 std=9.999614e-01 min=-3.843750e+00 max=4.312500e+00 l2_norm=3.135334e+02 nan_count=0 inf_count=0
  [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/gate_logits shape=(12, 128, 4) dtype=bfloat16 mean=-4.968616e-02 std=9.442701e-01 min=-3.515625e+00 max=3.187500e+00 l2_norm=7.411775e+01 nan_count=0 inf_count=0
  [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/router_weights shape=(12, 128, 2) dtype=bfloat16 mean=5.000076e-01 std=1.925698e-01 min=1.757812e-02 max=9.843750e-01 l2_norm=2.969752e+01 nan_count=0 inf_count=0
  [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/expert_outputs shape=(12, 128, 64) dtype=bfloat16 mean=-1.053208e-03 std=3.509205e-02 min=-1.699219e-01 max=1.699219e-01 l2_norm=1.100753e+01 nan_count=0 inf_count=0
  [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/expert_outputs/grad shape=(12, 128, 64) dtype=bfloat16 mean=-3.270777e-07 std=7.611285e-05 min=-2.393723e-04 max=2.355576e-04 l2_norm=2.386424e-02 nan_count=0 inf_count=0
  [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/router_weights/grad shape=(12, 128, 2) dtype=bfloat16 mean=-3.831090e-07 std=2.770081e-05 min=-1.196861e-04 max=1.049042e-04 l2_norm=1.535482e-03 nan_count=0 inf_count=0
  [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/gate_logits/grad shape=(12, 128, 4) dtype=bfloat16 mean=0.000000e+00 std=6.066302e-06 min=-4.196167e-05 max=4.196167e-05 l2_norm=4.754990e-04 nan_count=0 inf_count=0
  [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/expert_inputs/grad shape=(12, 128, 64) dtype=bfloat16 mean=-1.991233e-09 std=4.040187e-06 min=-2.288818e-05 max=2.342463e-05 l2_norm=1.266739e-03 nan_count=0 inf_count=0
  

[codecov]

Codecov Report

❌ Patch coverage is 88.16327% with 29 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/debug_tensor_interceptors.py	87.07%	9 Missing and 10 partials ⚠️
src/maxtext/layers/nnx_wrappers.py	55.55%	3 Missing and 1 partial ⚠️
src/maxtext/utils/debug_tensor_utils.py	94.66%	2 Missing and 2 partials ⚠️
src/maxtext/trainers/pre_train/train.py	85.71%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!