class_docstr: Standard Multi-Layer Perceptron; also known as a feed-forward network.
  
  !!! faq
  
      If you get a TypeError saying an object is not a valid JAX type, see the
          [FAQ](https://docs.kidger.site/equinox/faq/).
docstr: **Arguments**:
  
  - `in_size`: The input size. The input to the module should be a vector of
      shape `(in_features,)`
  - `out_size`: The output size. The output from the module will be a vector
      of shape `(out_features,)`.
  - `width_size`: The size of each hidden layer.
  - `depth`: The number of hidden layers, including the output layer.
      For example, `depth=2` results in an network with layers:
      [`Linear(in_size, width_size)`, `Linear(width_size, width_size)`,
      `Linear(width_size, out_size)`].
  - `activation`: The activation function after each hidden layer. Defaults to
      ReLU.
  - `final_activation`: The activation function after the output layer. Defaults
      to the identity.
  - `use_bias`: Whether to add on a bias to internal layers. Defaults
      to `True`.
  - `use_final_bias`: Whether to add on a bias to the final layer. Defaults
      to `True`.
  - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
      initialisation. (Keyword only argument.)
  
  Note that `in_size` also supports the string `"scalar"` as a special value.
  In this case the input to the module should be of shape `()`.
  
  Likewise `out_size` can also be a string `"scalar"`, in which case the
  output from the module will have shape `()`.
docstr_tok: ["*","*","Arguments","*","*",":","-","`","in_size","`",":","The","input","size",".","The","input","to","the","module","should","be","a","vector","of","shape","`","(","in_features",",",")","`","-","`","out_size","`",":","The","output","size",".","The","output","from","the","module","will","be","a","vector","of","shape","`","(","out_features",",",")","`",".","-","`","width_size","`",":","The","size","of","each","hidden","layer",".","-","`","depth","`",":","The","number","of","hidden","layers",",","including","the","output","layer",".","For","example",",","`","depth=2","`","results","in","an","network","with","layers",":","[","`","Linear","(","in_size",",","width_size",")","`",",","`","Linear","(","width_size",",","width_size",")","`",",","`","Linear","(","width_size",",","out_size",")","`","]",".","-","`","activation","`",":","The","activation","function","after","each","hidden","layer",".","Defaults","to","ReLU",".","-","`","final_activation","`",":","The","activation","function","after","the","output","layer",".","Defaults","to","the","identity",".","-","`","use_bias","`",":","Whether","to","add","on","a","bias","to","internal","layers",".","Defaults","to","`","True","`",".","-","`","use_final_bias","`",":","Whether","to","add","on","a","bias","to","the","final","layer",".","Defaults","to","`","True","`",".","-","`","key","`",":","A","`","jax.random.PRNGKey","`","used","to","provide","randomness","for","parameter","initialisation",".","(","Keyword","only","argument",".",")","Note","that","`","in_size","`","also","supports","the","string","`","``","scalar","''","`","as","a","special","value",".","In","this","case","the","input","to","the","module","should","be","of","shape","`","(",")","`",".","Likewise","`","out_size","`","can","also","be","a","string","`","``","scalar","''","`",",","in","which","case","the","output","from","the","module","will","have","shape","`","(",")","`","."]
code: def __init__(
      self,
      in_size: Union[int, Literal["scalar"]],
      out_size: Union[int, Literal["scalar"]],
      width_size: int,
      depth: int,
      activation: Callable = _relu,
      final_activation: Callable = _identity,
      use_bias: bool = True,
      use_final_bias: bool = True,
      *,
      key: PRNGKeyArray,
  ):
      """**Arguments**:
  
      - `in_size`: The input size. The input to the module should be a vector of
          shape `(in_features,)`
      - `out_size`: The output size. The output from the module will be a vector
          of shape `(out_features,)`.
      - `width_size`: The size of each hidden layer.
      - `depth`: The number of hidden layers, including the output layer.
          For example, `depth=2` results in an network with layers:
          [`Linear(in_size, width_size)`, `Linear(width_size, width_size)`,
          `Linear(width_size, out_size)`].
      - `activation`: The activation function after each hidden layer. Defaults to
          ReLU.
      - `final_activation`: The activation function after the output layer. Defaults
          to the identity.
      - `use_bias`: Whether to add on a bias to internal layers. Defaults
          to `True`.
      - `use_final_bias`: Whether to add on a bias to the final layer. Defaults
          to `True`.
      - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
          initialisation. (Keyword only argument.)
  
      Note that `in_size` also supports the string `"scalar"` as a special value.
      In this case the input to the module should be of shape `()`.
  
      Likewise `out_size` can also be a string `"scalar"`, in which case the
      output from the module will have shape `()`.
      """
  
      keys = jrandom.split(key, depth + 1)
      layers = []
      if depth == 0:
          layers.append(
              Linear(in_size, out_size, use_final_bias, key=keys[0])
          )
      else:
          layers.append(
              Linear(in_size, width_size, use_bias, key=keys[0])
          )
          for i in range(depth - 1):
              layers.append(
                  Linear(
                      width_size, width_size, use_bias, key=keys[i + 1]
                  )
              )
          layers.append(
              Linear(width_size, out_size, use_final_bias, key=keys[-1])
          )
      self.layers = tuple(layers)
      self.in_size = in_size
      self.out_size = out_size
      self.width_size = width_size
      self.depth = depth
      # In case `activation` or `final_activation` are learnt, then make a separate
      # copy of their weights for every neuron.
      self.activation = filter_vmap(
          filter_vmap(lambda: activation, axis_size=width_size),
          axis_size=depth,
      )()
      if out_size == "scalar":
          self.final_activation = final_activation
      else:
          self.final_activation = filter_vmap(
              lambda: final_activation, axis_size=out_size
          )()
      self.use_bias = use_bias
      self.use_final_bias = use_final_bias
  
file_path: equinox/equinox/nn/_mlp.py
imports: from collections.abc import Callable
  from typing import Literal, Optional, Union
  import jax
  import jax.nn
  import jax.random
  import jax.tree_util
  from jaxtyping import Array, PRNGKeyArray
  from .._doc_utils import doc_repr
  from .._filters import is_array
  from .._module import field, Module
  from .._vmap_pmap import filter_vmap
  from ._linear import Linear
label_desc: Use image node_id 1 to create a new MLP object from inherited base classes: Module with example: obj = MLP(in_size, out_size, width_size, depth, activation, final_activation, use_bias, use_final_bias)