class_docstr: A single step of a Gated Recurrent Unit (GRU).
  
  !!! example
  
      This is often used by wrapping it into a `jax.lax.scan`. For example:
  
      ```python
      class Model(Module):
          cell: GRUCell
  
          def __init__(self, **kwargs):
              self.cell = GRUCell(**kwargs)
  
          def __call__(self, xs):
              scan_fn = lambda state, input: (self.cell(input, state), None)
              init_state = jnp.zeros(self.cell.hidden_size)
              final_state, _ = jax.lax.scan(scan_fn, init_state, xs)
              return final_state
      ```
docstr: **Arguments:**
  
  - `input_size`: The dimensionality of the input vector at each time step.
  - `hidden_size`: The dimensionality of the hidden state passed along between
      time steps.
  - `use_bias`: Whether to add on a bias after each update.
  - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
      initialisation. (Keyword only argument.)
docstr_tok: ["*","*","Arguments",":","*","*","-","`","input_size","`",":","The","dimensionality","of","the","input","vector","at","each","time","step",".","-","`","hidden_size","`",":","The","dimensionality","of","the","hidden","state","passed","along","between","time","steps",".","-","`","use_bias","`",":","Whether","to","add","on","a","bias","after","each","update",".","-","`","key","`",":","A","`","jax.random.PRNGKey","`","used","to","provide","randomness","for","parameter","initialisation",".","(","Keyword","only","argument",".",")"]
code: def __init__(
      self,
      input_size: int,
      hidden_size: int,
      use_bias: bool = True,
      *,
      key: PRNGKeyArray,
  ):
      """**Arguments:**
  
      - `input_size`: The dimensionality of the input vector at each time step.
      - `hidden_size`: The dimensionality of the hidden state passed along between
          time steps.
      - `use_bias`: Whether to add on a bias after each update.
      - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
          initialisation. (Keyword only argument.)
      """
  
      ihkey, hhkey, bkey, bkey2 = jrandom.split(key, 4)
      lim = math.sqrt(1 / hidden_size)
  
      self.weight_ih = jrandom.uniform(
          ihkey, (3 * hidden_size, input_size), minval=-lim, maxval=lim
      )
      self.weight_hh = jrandom.uniform(
          hhkey, (3 * hidden_size, hidden_size), minval=-lim, maxval=lim
      )
      if use_bias:
          self.bias = jrandom.uniform(
              bkey, (3 * hidden_size,), minval=-lim, maxval=lim
          )
          self.bias_n = jrandom.uniform(
              bkey2, (hidden_size,), minval=-lim, maxval=lim
          )
      else:
          self.bias = None
          self.bias_n = None
  
      self.input_size = input_size
      self.hidden_size = hidden_size
      self.use_bias = use_bias
  
file_path: equinox/equinox/nn/_rnn.py
imports: import math
  from typing import Optional
  import jax
  import jax.nn
  import jax.numpy
  import jax.random
  from jaxtyping import Array, PRNGKeyArray
  from .._module import field, Module
label_desc: Use image node_id 1 to create a new GRUCell object from inherited base classes: Module with example: obj = GRUCell(input_size, hidden_size, use_bias)