code: def _get_norm_per_neuron(param, axes):
      return jnp.sqrt(jnp.sum(jnp.power(param, 2), axis=axes))
  
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 6 for calling a global function with example usage: _get_norm_per_neuron(param, axes) and returns: jnp