code: def weight_std(fn, kernel_name="kernel", eps=1e-8):
      def std(variables):
          params = variables["params"]
          assert kernel_name in params
          kernel = params[kernel_name]
          redux = tuple(range(kernel.ndim - 1))
          norm = jnp.square(kernel).sum(redux, keepdims=True)
          std_kernel = kernel / jnp.sqrt(norm + eps)
          params[kernel_name] = std_kernel
          return variables
  
      # map_variables handles a few of nasty edge cases here...
      # the transformed kind will be immutable inside fn
      # this way we avoid lost mutations to param
      # map_variables also avoids accidental reuse of rngs
      # and it makes sure that other state is updated correctly (not twice during init!)
      return lift.map_variables(fn, "params", std, init=True)
  
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 a global function with example usage: weight_std(fn, kernel_name, eps) and returns: lift, variables