code: def test_weight_std(self):
      x = random.normal(
          random.key(0),
          (
              1,
              4,
          ),
      )
      y, variables = init(mlp)(random.key(1), x)
  
      param_shapes = unfreeze(
          jax.tree_util.tree_map(jnp.shape, variables["params"])
      )
      self.assertEqual(
          param_shapes,
          {
              "hidden_0": {"kernel": (4, 8), "bias": (8,)},
              "out": {"kernel": (8, 1), "bias": (1,)},
          },
      )
      self.assertEqual(y.shape, (1, 1))
      self.assertTrue(y.ravel() < 1.0)
  
      y2 = apply(mlp)(variables, x)
      self.assertTrue(jnp.allclose(y, y2))
  
file_path: flax/tests/core/design/core_weight_std_test.py
imports: from functools import partial
  from typing import Sequence
  import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, apply, init, lift, nn, unfreeze
label_desc: Use image node_id 1 for calling the WeightStdTest obj's underlying member method code with example usage: obj.test_weight_std() without return types