code: def _l2_normalize(x, axis=None, eps=1e-12):
      """Normalizes along dimension `axis` using an L2 norm.
  
      This specialized function exists for numerical stability reasons.
  
      Args:
        x: An input ndarray.
        axis: Dimension along which to normalize, e.g. `1` to separately normalize
          vectors in a batch. Passing `None` views `t` as a flattened vector when
          calculating the norm (equivalent to Frobenius norm).
        eps: Epsilon to avoid dividing by zero.
  
      Returns:
        An array of the same shape as 'x' L2-normalized along 'axis'.
      """
      return x * jax.lax.rsqrt(
          (x * x).sum(axis=axis, keepdims=True) + eps
      )
  
imports: import re
  from typing import Optional
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import initializers
  from haiku._src import module
  import jax
  import jax.lax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _l2_normalize(x, axis, eps) and returns: unknown