class_docstr: LayerNorm module.
  
  See: https://arxiv.org/abs/1607.06450.
  
  Example usage:
  
  >>> ln = hk.LayerNorm(axis=-1, param_axis=-1,
  ...                   create_scale=True, create_offset=True)
  >>> x = ln(jnp.ones([8, 224, 224, 3]))
class_docstr_tok: ["LayerNorm","module",".","See",":","https",":","\/\/arxiv.org\/abs\/1607.06450",".","Example","usage",":",">",">",">","ln","=","hk.LayerNorm","(","axis=-1",",","param_axis=-1",",","...","create_scale=True",",","create_offset=True",")",">",">",">","x","=","ln","(","jnp.ones","(","[","8",",","224",",","224",",","3","]",")",")"]
docstr: Connects the layer norm.
  
  Args:
    inputs: An array, where the data format is ``[N, ..., C]``.
    scale: An array up to n-D. The shape of this tensor must be broadcastable
      to the shape of ``inputs``. This is the scale applied to the normalized
      inputs. This cannot be passed in if the module was constructed with
      ``create_scale=True``.
    offset: An array up to n-D. The shape of this tensor must be broadcastable
      to the shape of ``inputs``. This is the offset applied to the normalized
      inputs. This cannot be passed in if the module was constructed with
      ``create_offset=True``.
  
  Returns:
    The array, normalized.
docstr_tok: ["Connects","the","layer","norm",".","Args",":","inputs",":","An","array",",","where","the","data","format","is","``","[","N",",","...",",","C","]","``",".","scale",":","An","array","up","to","n-D",".","The","shape","of","this","tensor","must","be","broadcastable","to","the","shape","of","``","inputs","``",".","This","is","the","scale","applied","to","the","normalized","inputs",".","This","can","not","be","passed","in","if","the","module","was","constructed","with","``","create_scale=True","``",".","offset",":","An","array","up","to","n-D",".","The","shape","of","this","tensor","must","be","broadcastable","to","the","shape","of","``","inputs","``",".","This","is","the","offset","applied","to","the","normalized","inputs",".","This","can","not","be","passed","in","if","the","module","was","constructed","with","``","create_offset=True","``",".","Returns",":","The","array",",","normalized","."]
code: def __call__(
      self,
      inputs: jax.Array,
      scale: Optional[jax.Array] = None,
      offset: Optional[jax.Array] = None,
  ) -> jax.Array:
      """Connects the layer norm.
  
      Args:
        inputs: An array, where the data format is ``[N, ..., C]``.
        scale: An array up to n-D. The shape of this tensor must be broadcastable
          to the shape of ``inputs``. This is the scale applied to the normalized
          inputs. This cannot be passed in if the module was constructed with
          ``create_scale=True``.
        offset: An array up to n-D. The shape of this tensor must be broadcastable
          to the shape of ``inputs``. This is the offset applied to the normalized
          inputs. This cannot be passed in if the module was constructed with
          ``create_offset=True``.
  
      Returns:
        The array, normalized.
      """
      if self.create_scale and scale is not None:
          raise ValueError(
              "Cannot pass `scale` at call time if `create_scale=True`."
          )
      if self.create_offset and offset is not None:
          raise ValueError(
              "Cannot pass `offset` at call time if `create_offset=True`."
          )
  
      axis = to_abs_axes(self.axis, inputs.ndim)
      mean = jnp.mean(inputs, axis=axis, keepdims=True)
      if self.use_fast_variance:
          mean_of_squares = jnp.mean(
              jnp.square(inputs), axis=axis, keepdims=True
          )
          variance = mean_of_squares - jnp.square(mean)
      else:
          variance = jnp.var(inputs, axis=axis, keepdims=True)
  
      if (
          self.create_scale or self.create_offset
      ) and not self._param_axis_passed_explicitly:
          if ERROR_IF_PARAM_AXIS_NOT_EXPLICIT and axis != (
              inputs.ndim - 1,
          ):
              raise ValueError(
                  "When axis is not the final dimension we require "
                  "you to also pass `param_axis` in the ctor."
                  f" axis={axis} ndim={inputs.ndim}"
              )
  
      # Shape for the learnable scale and offset is the number of channels.
      # See: https://arxiv.org/pdf/1803.08494.pdf around equation 6.
      param_axis = to_abs_axes(self.param_axis, inputs.ndim)
      if param_axis == (inputs.ndim - 1,):
          # For param_axis=-1 we store non-broadcast param shape for compatibility
          # with older checkpoints.
          param_shape = (inputs.shape[-1],)
      else:
          param_shape = tuple(
              (inputs.shape[i] if i in param_axis else 1)
              for i in range(inputs.ndim)
          )
  
      if self.create_scale:
          scale = hk.get_parameter(
              "scale", param_shape, inputs.dtype, init=self.scale_init
          )
      elif scale is None:
          scale = np.array(1.0, dtype=inputs.dtype)
  
      if self.create_offset:
          offset = hk.get_parameter(
              "offset", param_shape, inputs.dtype, init=self.offset_init
          )
      elif offset is None:
          offset = np.array(0.0, dtype=inputs.dtype)
  
      if jax.config.jax_numpy_rank_promotion != "allow":
          # TODO(b/234327547): Explicit bcast triggers excessive mem usage on TPU.
          # We should remove the conditional (and always broadcast) when the
          # referenced bug is fixed.
          scale = jnp.broadcast_to(scale, inputs.shape)
          offset = jnp.broadcast_to(offset, inputs.shape)
          mean = jnp.broadcast_to(mean, inputs.shape)
  
      eps = jax.lax.convert_element_type(self.eps, variance.dtype)
      inv = scale * jax.lax.rsqrt(variance + eps)
      return inv * (inputs - mean) + offset
  
imports: import collections.abc
  from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the LayerNorm obj's underlying member method code with example usage: obj.__call__(inputs, scale, offset) and returns: unknown