Update L2Norm module to use NNX

[bvandermoon]

Description

Update the L2Norm module to use NNX instead of Linen.

L2Norm has no dependencies/params. And only one module depends on it (Attention - and only for Llama4).

Tests

Llama4 train_compile passes:

python3 -m pytest MaxText/tests/train_compile_test.py::TrainCompile::test_moe_llama4_17b_16e

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.