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`: The input, which should be a JAX array of shape `(input_size,)`.
  - `hidden`: The hidden state, which should be a JAX array of shape
      `(hidden_size,)`.
  - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
      (Keyword only argument.)
  
  **Returns:**
  
  The updated hidden state, which is a JAX array of shape `(hidden_size,)`.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","input","`",":","The","input",",","which","should","be","a","JAX","array","of","shape","`","(","input_size",",",")","`",".","-","`","hidden","`",":","The","hidden","state",",","which","should","be","a","JAX","array","of","shape","`","(","hidden_size",",",")","`",".","-","`","key","`",":","Ignored",";","provided","for","compatibility","with","the","rest","of","the","Equinox","API",".","(","Keyword","only","argument",".",")","*","*","Returns",":","*","*","The","updated","hidden","state",",","which","is","a","JAX","array","of","shape","`","(","hidden_size",",",")","`","."]
code: def __call__(
      self,
      input: Array,
      hidden: Array,
      *,
      key: Optional[PRNGKeyArray] = None,
  ):
      """**Arguments:**
  
      - `input`: The input, which should be a JAX array of shape `(input_size,)`.
      - `hidden`: The hidden state, which should be a JAX array of shape
          `(hidden_size,)`.
      - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
          (Keyword only argument.)
  
      **Returns:**
  
      The updated hidden state, which is a JAX array of shape `(hidden_size,)`.
      """
      if self.use_bias:
          bias = self.bias
          bias_n = self.bias_n
      else:
          bias = 0
          bias_n = 0
      igates = jnp.split(self.weight_ih @ input + bias, 3)
      hgates = jnp.split(self.weight_hh @ hidden, 3)
      reset = jnn.sigmoid(igates[0] + hgates[0])
      inp = jnn.sigmoid(igates[1] + hgates[1])
      new = jnn.tanh(igates[2] + reset * (hgates[2] + bias_n))
      return new + inp * (hidden - new)
  
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 2 for calling the GRUCell obj's underlying member method code with example usage: obj.__call__(input, hidden) and returns: unknown