Merge pull request #5076 from google:update-TestLayersSameGraph

PiperOrigin-RevId: 830557965

Author: Flax Authors

tests/nnx/nn/linear_test.py

@@ -256,6 +256,7 @@ class TestLayersSameGraph(parameterized.TestCase):
               ("kernel_init", "out_kernel_init", "bias_init", "out_bias_init"),
           ),
           (nnx.BatchNorm, (3,), ("scale_init", "bias_init")),
+          (nnx.LayerNorm, (3,), ("scale_init", "bias_init")),
           (nnx.RMSNorm, (3,), ("scale_init",)),
           (nnx.GroupNorm, (6, 3), ("scale_init", "bias_init")),
           (nnx.InstanceNorm, (6,), ("scale_init", "bias_init")),