class_docstr: GRU cell.
  
  The mathematical definition of the cell is as follows
  
  .. math::
  
      \begin{array}{ll}
      r = \sigma(W_{ir} x + b_{ir} + W_{hr} h) \\
      z = \sigma(W_{iz} x + b_{iz} + W_{hz} h) \\
      n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\
      h' = (1 - z) * n + z * h \\
      \end{array}
  
  where x is the input and h, is the output of the previous time step.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> x = jax.random.normal(jax.random.key(0), (2, 3))
    >>> layer = nn.GRUCell(features=4)
    >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)
    >>> variables = layer.init(jax.random.key(2), carry, x)
    >>> new_carry, out = layer.apply(variables, carry, x)
  
  Attributes:
    features: number of output features.
    gate_fn: activation function used for gates (default: sigmoid).
    activation_fn: activation function used for output and memory update
      (default: tanh).
    kernel_init: initializer function for the kernels that transform
      the input (default: lecun_normal).
    recurrent_kernel_init: initializer function for the kernels that transform
      the hidden state (default: initializers.orthogonal()).
    bias_init: initializer for the bias parameters (default: initializers.zeros_init())
    dtype: the dtype of the computation (default: None).
    param_dtype: the dtype passed to parameter initializers (default: float32).
docstr: Initialize the RNN cell carry.
  
  Args:
    rng: random number generator passed to the init_fn.
    input_shape: a tuple providing the shape of the input to the cell.
  
  Returns:
    An initialized carry for the given RNN cell.
docstr_tok: ["Initialize","the","RNN","cell","carry",".","Args",":","rng",":","random","number","generator","passed","to","the","init_fn",".","input_shape",":","a","tuple","providing","the","shape","of","the","input","to","the","cell",".","Returns",":","An","initialized","carry","for","the","given","RNN","cell","."]
code: def initialize_carry(
      self, rng: PRNGKey, input_shape: Tuple[int, ...]
  ):
      """Initialize the RNN cell carry.
  
      Args:
        rng: random number generator passed to the init_fn.
        input_shape: a tuple providing the shape of the input to the cell.
  
      Returns:
        An initialized carry for the given RNN cell.
      """
      batch_dims = input_shape[:-1]
      mem_shape = batch_dims + (self.features,)
      return self.carry_init(rng, mem_shape, self.param_dtype)
  
imports: from functools import partial
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import jax
  import numpy
  from absl import logging
  from jax import numpy
  from jax import random
  from typing_extensions import Protocol
  from flax.core import lift
  from flax.core.frozen_dict import FrozenDict
  from flax.linen import initializers, transforms
  from flax.linen.activation import sigmoid, tanh
  from flax.linen.dtypes import promote_dtype
  from flax.linen.linear import Conv, Dense, PrecisionLike, default_kernel_init
  from flax.linen.module import Module, compact, nowrap
label_desc: Use image node_id 2 for calling the GRUCell obj's underlying member method code with example usage: obj.initialize_carry(rng, input_shape) and returns: self