class_docstr: The ``RNN`` module takes any :class:`RNNCellBase` instance and applies it over a sequence
  using :func:`flax.linen.scan`.
  
  Example::
  
    >>> import jax.numpy as jnp
    >>> import jax
    >>> import flax.linen as nn
  
    >>> x = jnp.ones((10, 50, 32)) # (batch, time, features)
    >>> lstm = nn.RNN(nn.LSTMCell(64))
    >>> variables = lstm.init(jax.random.key(0), x)
    >>> y = lstm.apply(variables, x)
    >>> y.shape # (batch, time, cell_size)
    (10, 50, 64)
  
  As shown above, RNN uses the ``cell_size`` argument to set the ``size`` argument for the cell's
  ``initialize_carry`` method, in practice this is typically the number of hidden units you want
  for the cell. However, this may vary depending on the cell you are using, for example the
  :class:`ConvLSTMCell` requires a ``size`` argument of the form
  ``(kernel_height, kernel_width, features)``::
  
    >>> x = jnp.ones((10, 50, 32, 32, 3)) # (batch, time, height, width, features)
    >>> conv_lstm = nn.RNN(nn.ConvLSTMCell(64, kernel_size=(3, 3)))
    >>> y, variables = conv_lstm.init_with_output(jax.random.key(0), x)
    >>> y.shape # (batch, time, height, width, features)
    (10, 50, 32, 32, 64)
  
  By default RNN expect the time dimension after the batch dimension (``(*batch, time, *features)``),
  if you set ``time_major=True`` RNN will instead expect the time dimesion to be at the beginning
  (``(time, *batch, *features)``)::
  
    >>> x = jnp.ones((50, 10, 32)) # (time, batch, features)
    >>> lstm = nn.RNN(nn.LSTMCell(64), time_major=True)
    >>> variables = lstm.init(jax.random.key(0), x)
    >>> y = lstm.apply(variables, x)
    >>> y.shape # (time, batch, cell_size)
    (50, 10, 64)
  
  The output is an array of shape ``(*batch, time, *cell_size)`` by default (typically), however
  if you set ``return_carry=True`` it will instead return a tuple of the final carry and the output::
  
    >>> x = jnp.ones((10, 50, 32)) # (batch, time, features)
    >>> lstm = nn.RNN(nn.LSTMCell(64), return_carry=True)
    >>> variables = lstm.init(jax.random.key(0), x)
    >>> carry, y = lstm.apply(variables, x)
    >>> jax.tree_map(jnp.shape, carry) # ((batch, cell_size), (batch, cell_size))
    ((10, 64), (10, 64))
    >>> y.shape # (batch, time, cell_size)
    (10, 50, 64)
  
  To support variable length sequences, you can pass a ``seq_lengths`` which is an integer
  array of shape ``(*batch)`` where each element is the length of the sequence in the batch.
  For example::
  
    >>> seq_lengths = jnp.array([3, 2, 5])
  
  The output elements corresponding to padding elements are NOT zeroed out. If ``return_carry``
  is set to ``True`` the carry will be the state of the last valid element of each sequence.
  
  RNN also accepts some of the arguments of :func:`flax.linen.scan`, by default they are set to
  work with cells like :class:`LSTMCell` and :class:`GRUCell` but they can be overriden as needed.
  Overriding default values to scan looks like this::
  
    >>> lstm = nn.RNN(
    ...   nn.LSTMCell(64),
    ...   unroll=1, variable_axes={}, variable_broadcast='params',
    ...   variable_carry=False, split_rngs={'params': False})
  
  Attributes:
    cell: an instance of :class:`RNNCellBase`.
    time_major: if ``time_major=False`` (default) it will expect inputs with shape
      ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``.
    return_carry: if ``return_carry=False`` (default) only the output sequence is returned,
      else it will return a tuple of the final carry and the output sequence.
    reverse: if ``reverse=False`` (default) the sequence is processed from left to right and
      returned in the original order, else it will be processed from right to left, and
      returned in reverse order. If ``seq_lengths`` is passed, padding will always remain
      at the end of the sequence.
    keep_order: if ``keep_order=True``, when ``reverse=True``
      the output will be reversed back to the original order after processing, this is
      useful to align sequences in bidirectional RNNs. If ``keep_order=False`` (default),
      the output will remain in the order specified by ``reverse``.
    unroll: how many scan iterations to unroll within a single iteration of a loop,
      defaults to 1. This argument will be passed to ``nn.scan``.
    variable_axes: a dictionary mapping each collection to either an integer ``i`` (meaning we scan over
      dimension ``i``) or ``None`` (replicate rather than scan). This argument is forwarded to ``nn.scan``.
    variable_broadcast: Specifies the broadcasted variable collections. A
      broadcasted variable should not depend on any computation that cannot be
      lifted out of the loop. This is typically used to define shared parameters
      inside the fn. This argument is forwarded to ``nn.scan``.
    variable_carry: Specifies the variable collections that are carried through
      the loop. Mutations to these variables are carried to the next iteration
      and will be preserved when the scan finishes. This argument is forwarded to
      ``nn.scan``.
    split_rngs: a mapping from PRNGSequenceFilter to bool specifying whether a collection's
      PRNG key should be split such that its values are different at each step, or replicated
      such that its values remain the same at each step. This argument is forwarded to ``nn.scan``.
docstr: Applies the RNN to the inputs.
  
  ``__call__`` allows you to optionally override some attributes like ``return_carry``
  and ``time_major`` defined in the constructor.
  
  Arguments:
    inputs: the input sequence.
    initial_carry: the initial carry, if not provided it will be initialized
      using the cell's :meth:`RNNCellBase.initialize_carry` method.
    init_key: a PRNG key used to initialize the carry, if not provided
      ``jax.random.key(0)`` will be used. Most cells will ignore this
      argument.
    seq_lengths: an optional integer array of shape ``(*batch)`` indicating
      the length of each sequence, elements whose index in the time dimension
      is greater than the corresponding length will be considered padding and
      will be ignored.
    return_carry: if ``return_carry=False`` (default) only the output sequence is returned,
      else it will return a tuple of the final carry and the output sequence.
    time_major: if ``time_major=False`` (default) it will expect inputs with shape
      ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``.
    reverse: overrides the ``reverse`` attribute, if ``reverse=False`` (default) the sequence is
      processed from left to right and returned in the original order, else it will be processed
      from right to left, and returned in reverse order. If ``seq_lengths`` is passed,
      padding will always remain at the end of the sequence.
    keep_order: overrides the ``keep_order`` attribute, if ``keep_order=True``, when ``reverse=True``
      the output will be reversed back to the original order after processing, this is
      useful to align sequences in bidirectional RNNs. If ``keep_order=False`` (default),
      the output will remain in the order specified by ``reverse``.
  Returns:
    if ``return_carry=False`` (default) only the output sequence is returned,
    else it will return a tuple of the final carry and the output sequence.
docstr_tok: ["Applies","the","RNN","to","the","inputs",".","``","__call__","``","allows","you","to","optionally","override","some","attributes","like","``","return_carry","``","and","``","time_major","``","defined","in","the","constructor",".","Arguments",":","inputs",":","the","input","sequence",".","initial_carry",":","the","initial","carry",",","if","not","provided","it","will","be","initialized","using","the","cell","'s",":","meth",":","`","RNNCellBase.initialize_carry","`","method",".","init_key",":","a","PRNG","key","used","to","initialize","the","carry",",","if","not","provided","``","jax.random.key","(","0",")","``","will","be","used",".","Most","cells","will","ignore","this","argument",".","seq_lengths",":","an","optional","integer","array","of","shape","``","(","*","batch",")","``","indicating","the","length","of","each","sequence",",","elements","whose","index","in","the","time","dimension","is","greater","than","the","corresponding","length","will","be","considered","padding","and","will","be","ignored",".","return_carry",":","if","``","return_carry=False","``","(","default",")","only","the","output","sequence","is","returned",",","else","it","will","return","a","tuple","of","the","final","carry","and","the","output","sequence",".","time_major",":","if","``","time_major=False","``","(","default",")","it","will","expect","inputs","with","shape","``","(","*","batch",",","time",",","*","features",")","``",",","else","it","will","expect","inputs","with","shape","``","(","time",",","*","batch",",","*","features",")","``",".","reverse",":","overrides","the","``","reverse","``","attribute",",","if","``","reverse=False","``","(","default",")","the","sequence","is","processed","from","left","to","right","and","returned","in","the","original","order",",","else","it","will","be","processed","from","right","to","left",",","and","returned","in","reverse","order",".","If","``","seq_lengths","``","is","passed",",","padding","will","always","remain","at","the","end","of","the","sequence",".","keep_order",":","overrides","the","``","keep_order","``","attribute",",","if","``","keep_order=True","``",",","when","``","reverse=True","``","the","output","will","be","reversed","back","to","the","original","order","after","processing",",","this","is","useful","to","align","sequences","in","bidirectional","RNNs",".","If","``","keep_order=False","``","(","default",")",",","the","output","will","remain","in","the","order","specified","by","``","reverse","``",".","Returns",":","if","``","return_carry=False","``","(","default",")","only","the","output","sequence","is","returned",",","else","it","will","return","a","tuple","of","the","final","carry","and","the","output","sequence","."]
code: def __call__(
      self,
      inputs: jax.Array,
      *,
      initial_carry: Optional[Carry] = None,
      init_key: Optional[PRNGKey] = None,
      seq_lengths: Optional[Array] = None,
      return_carry: Optional[bool] = None,
      time_major: Optional[bool] = None,
      reverse: Optional[bool] = None,
      keep_order: Optional[bool] = None,
  ) -> Union[Output, Tuple[Carry, Output]]:
      """
      Applies the RNN to the inputs.
  
      ``__call__`` allows you to optionally override some attributes like ``return_carry``
      and ``time_major`` defined in the constructor.
  
      Arguments:
        inputs: the input sequence.
        initial_carry: the initial carry, if not provided it will be initialized
          using the cell's :meth:`RNNCellBase.initialize_carry` method.
        init_key: a PRNG key used to initialize the carry, if not provided
          ``jax.random.key(0)`` will be used. Most cells will ignore this
          argument.
        seq_lengths: an optional integer array of shape ``(*batch)`` indicating
          the length of each sequence, elements whose index in the time dimension
          is greater than the corresponding length will be considered padding and
          will be ignored.
        return_carry: if ``return_carry=False`` (default) only the output sequence is returned,
          else it will return a tuple of the final carry and the output sequence.
        time_major: if ``time_major=False`` (default) it will expect inputs with shape
          ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``.
        reverse: overrides the ``reverse`` attribute, if ``reverse=False`` (default) the sequence is
          processed from left to right and returned in the original order, else it will be processed
          from right to left, and returned in reverse order. If ``seq_lengths`` is passed,
          padding will always remain at the end of the sequence.
        keep_order: overrides the ``keep_order`` attribute, if ``keep_order=True``, when ``reverse=True``
          the output will be reversed back to the original order after processing, this is
          useful to align sequences in bidirectional RNNs. If ``keep_order=False`` (default),
          the output will remain in the order specified by ``reverse``.
      Returns:
        if ``return_carry=False`` (default) only the output sequence is returned,
        else it will return a tuple of the final carry and the output sequence.
      """
  
      if return_carry is None:
          return_carry = self.return_carry
      if time_major is None:
          time_major = self.time_major
      if reverse is None:
          reverse = self.reverse
      if keep_order is None:
          keep_order = self.keep_order
  
      # Infer the number of batch dimensions from the input shape.
      # Cells like ConvLSTM have additional spatial dimensions.
      time_axis = (
          0
          if time_major
          else inputs.ndim - (self.cell.num_feature_axes + 1)
      )
  
      # make time_axis positive
      if time_axis < 0:
          time_axis += inputs.ndim
  
      if time_major:
          # we add +1 because we moved the time axis to the front
          batch_dims = inputs.shape[1 : -self.cell.num_feature_axes]
      else:
          batch_dims = inputs.shape[:time_axis]
  
      # maybe reverse the sequence
      if reverse:
          inputs = jax.tree_map(
              lambda x: flip_sequences(
                  x,
                  seq_lengths,
                  num_batch_dims=len(batch_dims),
                  time_major=time_major,  # type: ignore
              ),
              inputs,
          )
  
      carry: Carry
      if initial_carry is None:
          if init_key is None:
              init_key = random.key(0)
  
          input_shape = (
              inputs.shape[:time_axis] + inputs.shape[time_axis + 1 :]
          )
          carry = self.cell.initialize_carry(init_key, input_shape)
      else:
          carry = initial_carry
  
      slice_carry = seq_lengths is not None and return_carry
  
      def scan_fn(
          cell: RNNCellBase, carry: Carry, x: Array
      ) -> Union[
          Tuple[Carry, Array], Tuple[Carry, Tuple[Carry, Array]]
      ]:
          carry, y = cell(carry, x)
          # When we have a segmentation mask we return the carry as an output
          # so that we can select the last carry for each sequence later.
          # This uses more memory but is faster than using jnp.where at each
          # iteration. As a small optimization do this when we really need it.
          if slice_carry:
              return carry, (carry, y)
          else:
              return carry, y
  
      scan = transforms.scan(
          scan_fn,
          in_axes=time_axis,
          out_axes=(0, time_axis) if slice_carry else time_axis,
          unroll=self.unroll,
          variable_axes=self.variable_axes,
          variable_broadcast=self.variable_broadcast,
          variable_carry=self.variable_carry,
          split_rngs=self.split_rngs,
      )
  
      scan_output = scan(self.cell, carry, inputs)
  
      # Next we select the final carry. If a segmentation mask was provided and
      # return_carry is True we slice the carry history and select the last valid
      # carry for each sequence. Otherwise we just use the last carry.
      if slice_carry:
          assert seq_lengths is not None
          _, (carries, outputs) = scan_output
          # seq_lengths[None] expands the shape of the mask to match the
          # number of dimensions of the carry.
          carry = _select_last_carry(carries, seq_lengths)
      else:
          carry, outputs = scan_output
  
      if reverse and keep_order:
          outputs = jax.tree_map(
              lambda x: flip_sequences(
                  x,
                  seq_lengths,
                  num_batch_dims=len(batch_dims),
                  time_major=time_major,  # type: ignore
              ),
              outputs,
          )
  
      if return_carry:
          return carry, outputs
      else:
          return outputs
  
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 RNN obj's underlying member method code with example usage: obj.__call__(inputs) and returns: carry,  outputs, outputs, carry,  carry,  y, carry,  y