code: def norm_kernel_init_fn(rng, shape, dtype):
      """Initialize kernel with l2 normalized columns."""
      param = nn.linear.default_kernel_init(rng, shape, dtype)
      param /= jnp.linalg.norm(param, axis=0, keepdims=True) + 1e-10
      return param
  
imports: import copy
  import functools
  from typing import Any, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 2 for calling a global function with example usage: norm_kernel_init_fn(rng, shape, dtype) and returns: param