class_docstr: More efficient LSTM Cell that concatenates state components before matmul.
  
  The parameters are compatible with ``LSTMCell``. Note that this cell is often
  faster than ``LSTMCell`` as long as the hidden size is roughly <= 2048 units.
  
  The mathematical definition of the cell is the same as ``LSTMCell`` and 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.OptimizedLSTMCell(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:
    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: An optimized 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: ["An","optimized","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: Tuple[Array, Array], inputs: Array
  ) -> Tuple[Tuple[Array, Array], Array]:
      r"""An optimized 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]
  
      def _concat_dense(
          inputs: Array,
          params: Mapping[str, Tuple[Array, Optional[Array]]],
          use_bias: bool = True,
      ) -> Dict[str, Array]:
          # Concatenates the individual kernels and biases, given in params, into a
          # single kernel and single bias for efficiency before applying them using
          # dot_general.
          kernels = [kernel for kernel, _ in params.values()]
          kernel = jnp.concatenate(kernels, axis=-1)
          if use_bias:
              biases = []
              for _, bias in params.values():
                  if bias is None:
                      raise ValueError(
                          "bias is None but use_bias is True."
                      )
                  biases.append(bias)
              bias = jnp.concatenate(biases, axis=-1)
          else:
              bias = None
          inputs, kernel, bias = promote_dtype(
              inputs, kernel, bias, dtype=self.dtype
          )
          y = jnp.dot(inputs, kernel)
          if use_bias:
              # This assert is here since mypy can't infer that bias cannot be None
              assert bias is not None
              y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
  
          # Split the result back into individual (i, f, g, o) outputs.
          split_indices = np.cumsum(
              [kernel.shape[-1] for kernel in kernels[:-1]]
          )
          ys = jnp.split(y, split_indices, axis=-1)
          return dict(zip(params.keys(), ys))
  
      # Create params with the same names/shapes as `LSTMCell` for compatibility.
      dense_params_h = {}
      dense_params_i = {}
      for component in ["i", "f", "g", "o"]:
          dense_params_i[component] = DenseParams(
              features=hidden_features,
              use_bias=False,
              param_dtype=self.param_dtype,
              kernel_init=self.kernel_init,
              bias_init=self.bias_init,
              name=f"i{component}",  # type: ignore[call-arg]
          )(inputs)
          dense_params_h[component] = DenseParams(
              features=hidden_features,
              use_bias=True,
              param_dtype=self.param_dtype,
              kernel_init=self.recurrent_kernel_init,
              bias_init=self.bias_init,
              name=f"h{component}",  # type: ignore[call-arg]
          )(h)
      dense_h = _concat_dense(h, dense_params_h, use_bias=True)
      dense_i = _concat_dense(inputs, dense_params_i, use_bias=False)
  
      i = self.gate_fn(dense_h["i"] + dense_i["i"])
      f = self.gate_fn(dense_h["f"] + dense_i["f"])
      g = self.activation_fn(dense_h["g"] + dense_i["g"])
      o = self.gate_fn(dense_h["o"] + dense_i["o"])
  
      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 OptimizedLSTMCell obj's underlying member method code with example usage: obj.__call__(carry, inputs) and returns: new_c,  new_h,  new_h, dict