class_docstr: Performs a linear transformation.
docstr: **Arguments:**
  
  - `x`: The input. Should be a JAX array of shape `(in_features,)`. (Or shape
      `()` if `in_features="scalar"`.)
  - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
      (Keyword only argument.)
  
  !!! info
  
      If you want to use higher order tensors as inputs (for example featuring "
      "batch dimensions) then use `jax.vmap`. For example, for an input `x` of "
      "shape `(batch, in_features)`, using
      ```python
      linear = equinox.nn.Linear(...)
      jax.vmap(linear)(x)
      ```
      will produce the appropriate output of shape `(batch, out_features)`.
  
  **Returns:**
  
  A JAX array of shape `(out_features,)`. (Or shape `()` if
  `out_features="scalar"`.)
docstr_tok: ["*","*","Arguments",":","*","*","-","`","x","`",":","The","input",".","Should","be","a","JAX","array","of","shape","`","(","in_features",",",")","`",".","(","Or","shape","`","(",")","`","if","`","in_features=","''","scalar","''","`",".",")","-","`","key","`",":","Ignored",";","provided","for","compatibility","with","the","rest","of","the","Equinox","API",".","(","Keyword","only","argument",".",")","!","!","!","info","If","you","want","to","use","higher","order","tensors","as","inputs","(","for","example","featuring","``","``","batch","dimensions",")","then","use","`","jax.vmap","`",".","For","example",",","for","an","input","`","x","`","of","``","``","shape","`","(","batch",",","in_features",")","`",",","using","``","`","python","linear","=","equinox.nn.Linear","(","...",")","jax.vmap","(","linear",")","(","x",")","``","`","will","produce","the","appropriate","output","of","shape","`","(","batch",",","out_features",")","`",".","*","*","Returns",":","*","*","A","JAX","array","of","shape","`","(","out_features",",",")","`",".","(","Or","shape","`","(",")","`","if","`","out_features=","''","scalar","''","`",".",")"]
code: def __call__(
      self, x: Array, *, key: Optional[PRNGKeyArray] = None
  ) -> Array:
      """**Arguments:**
  
      - `x`: The input. Should be a JAX array of shape `(in_features,)`. (Or shape
          `()` if `in_features="scalar"`.)
      - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
          (Keyword only argument.)
  
      !!! info
  
          If you want to use higher order tensors as inputs (for example featuring "
          "batch dimensions) then use `jax.vmap`. For example, for an input `x` of "
          "shape `(batch, in_features)`, using
          ```python
          linear = equinox.nn.Linear(...)
          jax.vmap(linear)(x)
          ```
          will produce the appropriate output of shape `(batch, out_features)`.
  
      **Returns:**
  
      A JAX array of shape `(out_features,)`. (Or shape `()` if
      `out_features="scalar"`.)
      """
  
      if self.in_features == "scalar":
          if jnp.shape(x) != ():
              raise ValueError("x must have scalar shape")
          x = jnp.broadcast_to(x, (1,))
      x = self.weight @ x
      if self.bias is not None:
          x = x + self.bias
      if self.out_features == "scalar":
          assert jnp.shape(x) == (1,)
          x = jnp.squeeze(x)
      return x
  
file_path: equinox/equinox/nn/_linear.py
imports: import math
  from typing import Any, Literal, Optional, TypeVar, Union
  import jax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, PRNGKeyArray
  from .._module import field, Module
label_desc: Use image node_id 2 for calling the Linear obj's underlying member method code with example usage: obj.__call__(x) and returns: x