class_docstr: LSTM cell.
  
  The mathematical definition of the cell is as follows
  
  .. math::
      \begin{array}{ll}
      i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\
      f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\
      g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\
      o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\
      c' = f * c + i * g \\
      h' = o * \tanh(c') \\
      \end{array}
  
  where x is the input, h is the output of the previous time step, and c is
  the memory.
  
  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.LSTMCell(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: infer from inputs and params).
    param_dtype: the dtype passed to parameter initializers (default: float32).
docstr: A long short-term memory (LSTM) cell.
  
  Args:
    carry: the hidden state of the LSTM cell,
      initialized using ``LSTMCell.initialize_carry``.
    inputs: an ndarray with the input for the current time step.
      All dimensions except the final are considered batch dimensions.
  
  Returns:
    A tuple with the new carry and the output.
docstr_tok: ["A","long","short-term","memory","(","LSTM",")","cell",".","Args",":","carry",":","the","hidden","state","of","the","LSTM","cell",",","initialized","using","``","LSTMCell.initialize_carry","``",".","inputs",":","an","ndarray","with","the","input","for","the","current","time","step",".","All","dimensions","except","the","final","are","considered","batch","dimensions",".","Returns",":","A","tuple","with","the","new","carry","and","the","output","."]
code: def __call__(self, carry, inputs):
      r"""A long short-term memory (LSTM) cell.
  
      Args:
        carry: the hidden state of the LSTM cell,
          initialized using ``LSTMCell.initialize_carry``.
        inputs: an ndarray with the input for the current time step.
          All dimensions except the final are considered batch dimensions.
  
      Returns:
        A tuple with the new carry and the output.
      """
      c, h = carry
      hidden_features = h.shape[-1]
      # input and recurrent layers are summed so only one needs a bias.
      dense_h = partial(
          Dense,
          features=hidden_features,
          use_bias=True,
          kernel_init=self.recurrent_kernel_init,
          bias_init=self.bias_init,
          dtype=self.dtype,
          param_dtype=self.param_dtype,
      )
      dense_i = partial(
          Dense,
          features=hidden_features,
          use_bias=False,
          kernel_init=self.kernel_init,
          dtype=self.dtype,
          param_dtype=self.param_dtype,
      )
      i = self.gate_fn(
          dense_i(name="ii")(inputs) + dense_h(name="hi")(h)
      )
      f = self.gate_fn(
          dense_i(name="if")(inputs) + dense_h(name="hf")(h)
      )
      g = self.activation_fn(
          dense_i(name="ig")(inputs) + dense_h(name="hg")(h)
      )
      o = self.gate_fn(
          dense_i(name="io")(inputs) + dense_h(name="ho")(h)
      )
      new_c = f * c + i * g
      new_h = o * self.activation_fn(new_c)
      return (new_c, new_h), new_h
  
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 1 for calling the LSTMCell obj's underlying member method code with example usage: obj.__call__(carry, inputs) and returns: new_c,  new_h,  new_h