code: def _weight_normalization_per_neuron_norm(param, axes):
      norm_per_neuron = _get_norm_per_neuron(param, axes)
      norm_per_neuron = jnp.expand_dims(norm_per_neuron, axis=axes)
      normalized_param = param / norm_per_neuron
      return normalized_param
  
imports: import functools
  import logging
  import flax
  from flax import linen
  import gin
  import jax
  from jax import random
  import jax.numpy
  import optax
label_desc: Use image node_id 5 for calling a global function with example usage: _weight_normalization_per_neuron_norm(param, axes) and returns: normalized_param