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: Constructs a LayerNorm module.
  
  Args:
    axis: Integer, list of integers, or slice indicating which axes to
      normalize over. Note that the shape of the scale/offset parameters are
      controlled by the ``param_axis`` argument.
    create_scale: Bool, defines whether to create a trainable scale
      per channel applied after the normalization.
    create_offset: Bool, defines whether to create a trainable offset
      per channel applied after normalization and scaling.
    eps: Small epsilon to avoid division by zero variance. Defaults ``1e-5``,
      as in the paper and Sonnet.
    scale_init: Optional initializer for gain (aka scale). By default, one.
    offset_init: Optional initializer for bias (aka offset). By default, zero.
    use_fast_variance: If true, use a faster but less numerically stable
      formulation for computing variance.
    name: The module name.
    param_axis: Axis used to determine the parameter shape of the learnable
      scale/offset. Sonnet sets this to the channel/feature axis (e.g. to
      ``-1`` for ``NHWC``). Other libraries set this to the same as the
      reduction axis (e.g. ``axis=param_axis``).
docstr_tok: ["Constructs","a","LayerNorm","module",".","Args",":","axis",":","Integer",",","list","of","integers",",","or","slice","indicating","which","axes","to","normalize","over",".","Note","that","the","shape","of","the","scale\/offset","parameters","are","controlled","by","the","``","param_axis","``","argument",".","create_scale",":","Bool",",","defines","whether","to","create","a","trainable","scale","per","channel","applied","after","the","normalization",".","create_offset",":","Bool",",","defines","whether","to","create","a","trainable","offset","per","channel","applied","after","normalization","and","scaling",".","eps",":","Small","epsilon","to","avoid","division","by","zero","variance",".","Defaults","``","1e-5","``",",","as","in","the","paper","and","Sonnet",".","scale_init",":","Optional","initializer","for","gain","(","aka","scale",")",".","By","default",",","one",".","offset_init",":","Optional","initializer","for","bias","(","aka","offset",")",".","By","default",",","zero",".","use_fast_variance",":","If","true",",","use","a","faster","but","less","numerically","stable","formulation","for","computing","variance",".","name",":","The","module","name",".","param_axis",":","Axis","used","to","determine","the","parameter","shape","of","the","learnable","scale\/offset",".","Sonnet","sets","this","to","the","channel\/feature","axis","(","e.g",".","to","``","-1","``","for","``","NHWC","``",")",".","Other","libraries","set","this","to","the","same","as","the","reduction","axis","(","e.g",".","``","axis=param_axis","``",")","."]
code: def __init__(
      self,
      axis: AxisOrAxes,
      create_scale: bool,
      create_offset: bool,
      eps: float = 1e-5,
      scale_init: Optional[hk.initializers.Initializer] = None,
      offset_init: Optional[hk.initializers.Initializer] = None,
      use_fast_variance: bool = False,
      name: Optional[str] = None,
      *,
      param_axis: Optional[AxisOrAxes] = None,
  ):
      """Constructs a LayerNorm module.
  
      Args:
        axis: Integer, list of integers, or slice indicating which axes to
          normalize over. Note that the shape of the scale/offset parameters are
          controlled by the ``param_axis`` argument.
        create_scale: Bool, defines whether to create a trainable scale
          per channel applied after the normalization.
        create_offset: Bool, defines whether to create a trainable offset
          per channel applied after normalization and scaling.
        eps: Small epsilon to avoid division by zero variance. Defaults ``1e-5``,
          as in the paper and Sonnet.
        scale_init: Optional initializer for gain (aka scale). By default, one.
        offset_init: Optional initializer for bias (aka offset). By default, zero.
        use_fast_variance: If true, use a faster but less numerically stable
          formulation for computing variance.
        name: The module name.
        param_axis: Axis used to determine the parameter shape of the learnable
          scale/offset. Sonnet sets this to the channel/feature axis (e.g. to
          ``-1`` for ``NHWC``). Other libraries set this to the same as the
          reduction axis (e.g. ``axis=param_axis``).
      """
      super().__init__(name=name)
      if not create_scale and scale_init is not None:
          raise ValueError(
              "Cannot set `scale_init` if `create_scale=False`."
          )
      if not create_offset and offset_init is not None:
          raise ValueError(
              "Cannot set `offset_init` if `create_offset=False`."
          )
  
      self.axis = to_axes_or_slice(axis)
      self.eps = eps
      self.create_scale = create_scale
      self.create_offset = create_offset
      self.scale_init = scale_init or jnp.ones
      self.offset_init = offset_init or jnp.zeros
      self.use_fast_variance = use_fast_variance
      self._param_axis_passed_explicitly = param_axis is not None
      self.param_axis = (
          (-1,) if param_axis is None else to_axes_or_slice(param_axis)
      )
  
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 1 to create a new LayerNorm object from inherited base classes: hk with example: obj = LayerNorm(axis, create_scale, create_offset, eps, scale_init, offset_init, use_fast_variance, name)