code: def weight_standardize(
      w: jnp.ndarray, axis: Union[Sequence[int], int], eps: float
  ):
      """Standardize (mean=0, var=1) a weight."""
      w = w - jnp.mean(w, axis=axis, keepdims=True)
      w = w / jnp.sqrt(
          jnp.mean(jnp.square(w), axis=axis, keepdims=True) + eps
      )
      return w
  
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling a global function with example usage: weight_standardize(w, axis, eps) and returns: w