base class for recurrent layers args cell a rnn cell instance or a list of rnn cell instances a rnn cell is a class that has a callinputatt statesatt method returning outputatt statesattplus1 the call method of the cell can also take the optional argument constants see section note on passing external constants below a statesize attribute this can be a single integer single state in which case it is the size of the recurrent state this can also be a listtuple of integers one size per state a outputsize attribute a single integer a getinitialstatebatchsizenone method that creates a tensor meant to be fed to call as the initial state if the user didn t specify any initial state via other means the returned initial state should have shape batchsize cell statesize the cell might choose to create a tensor full of zeros or other values based on the cell s implementation inputs is the input tensor to the rnn layer with shape batchsize timesteps features if this method is not implemented by the cell the rnn layer will create a zero filled tensor with shape batchsize cell statesize in the case that cell is a list of rnn cell instances the cells will be stacked on top of each other in the rnn resulting in an efficient stacked rnn returnsequences boolean default false whether to return the last output in the output sequence or the full sequence returnstate boolean default false whether to return the last state in addition to the output gobackwards boolean default false if true process the input sequence backwards and return the reversed sequence stateful boolean default false if true the last state for each sample at index i in a batch will be used as initial state for the sample of index i in the following batch unroll boolean default false if true the network will be unrolled else a symbolic loop will be used unrolling can speedup a rnn although it tends to be more memoryintensive unrolling is only suitable for short sequences zerooutputformask boolean default false whether the output should use zeros for the masked timesteps note that this field is only used when returnsequences is true and mask is provided it can useful if you want to reuse the raw output sequence of the rnn without interference from the masked timesteps e g merging bidirectional rnns call arguments inputs input tensor initialstate list of initial state tensors to be passed to the first call of the cell mask binary tensor of shape batchsize timesteps indicating whether a given timestep should be masked an individual true entry indicates that the corresponding timestep should be utilized while a false entry indicates that the corresponding timestep should be ignored training python boolean indicating whether the layer should behave in training mode or in inference mode this argument is passed to the cell when calling it this is for use with cells that use dropout input shape 3d tensor with shape batchsize timesteps features output shape if returnstate a list of tensors the first tensor is the output the remaining tensors are the last states each with shape batchsize statesize where statesize could be a high dimension tensor shape if returnsequences 3d tensor with shape batchsize timesteps outputsize masking this layer supports masking for input data with a variable number of timesteps to introduce masks to your data use a keras layers embedding layer with the maskzero parameter set to true note on using statefulness in rnns you can set rnn layers to be stateful which means that the states computed for the samples in one batch will be reused as initial states for the samples in the next batch this assumes a onetoone mapping between samples in different successive batches to enable statefulness specify statefultrue in the layer constructor specify a fixed batch size for your model by passing if sequential model batchinputshape to the first layer in your model else for functional model with 1 or more input layers batchshape to all the first layers in your model this is the expected shape of your inputs including the batch size it should be a tuple of integers e g 32 10 100 specify shufflefalse when calling fit to reset the states of your model call resetstates on either a specific layer or on your entire model note on specifying the initial state of rnns you can specify the initial state of rnn layers symbolically by calling them with the keyword argument initialstate the value of initialstate should be a tensor or list of tensors representing the initial state of the rnn layer you can specify the initial state of rnn layers numerically by calling resetstates with the keyword argument states the value of states should be a numpy array or list of numpy arrays representing the initial state of the rnn layer examples python from keras layers import rnn from keras import ops first let s define a rnn cell as a layer subclass class minimalrnncellkeras layers layer def initself units kwargs super initkwargs self units units self statesize units def buildself inputshape self kernel self addweightshapeinputshape1 self units initializer uniform name kernel self recurrentkernel self addweight shapeself units self units initializer uniform name recurrentkernel self built true def callself inputs states prevoutput states0 h ops matmulinputs self kernel output h ops matmulprevoutput self recurrentkernel return output output let s use this cell in a rnn layer cell minimalrnncell32 x keras inputnone 5 layer rnncell y layerx here s how to use the cell to build a stacked rnn cells minimalrnncell32 minimalrnncell64 x keras inputnone 5 layer rnncells y layerx if true the output for masked timestep will be zeros whereas in the false case output from previous timestep is returned for masked timestep time step masks must be the same for each input this is because the mask for an rnn is of size batch timesteps 1 and specifies which time steps should be skipped and a time step must be skipped for all inputs build cell if layer rnn expect the states in a list even if single state force the state to be a list in case it is a namedtuple eg lstmstatetuple compatibility alias rnn expect the states in a list even if single state cast states to compute dtype note that states may be deeply nested e g in the stacked cells case prepopulate the dropout state so that the innerloop is stateless this is particularly important for jax backend base class for recurrent layers args cell a rnn cell instance or a list of rnn cell instances a rnn cell is a class that has a call input_at_t states_at_t method returning output_at_t states_at_t_plus_1 the call method of the cell can also take the optional argument constants see section note on passing external constants below a state_size attribute this can be a single integer single state in which case it is the size of the recurrent state this can also be a list tuple of integers one size per state a output_size attribute a single integer a get_initial_state batch_size none method that creates a tensor meant to be fed to call as the initial state if the user didn t specify any initial state via other means the returned initial state should have shape batch_size cell state_size the cell might choose to create a tensor full of zeros or other values based on the cell s implementation inputs is the input tensor to the rnn layer with shape batch_size timesteps features if this method is not implemented by the cell the rnn layer will create a zero filled tensor with shape batch_size cell state_size in the case that cell is a list of rnn cell instances the cells will be stacked on top of each other in the rnn resulting in an efficient stacked rnn return_sequences boolean default false whether to return the last output in the output sequence or the full sequence return_state boolean default false whether to return the last state in addition to the output go_backwards boolean default false if true process the input sequence backwards and return the reversed sequence stateful boolean default false if true the last state for each sample at index i in a batch will be used as initial state for the sample of index i in the following batch unroll boolean default false if true the network will be unrolled else a symbolic loop will be used unrolling can speed up a rnn although it tends to be more memory intensive unrolling is only suitable for short sequences zero_output_for_mask boolean default false whether the output should use zeros for the masked timesteps note that this field is only used when return_sequences is true and mask is provided it can useful if you want to reuse the raw output sequence of the rnn without interference from the masked timesteps e g merging bidirectional rnns call arguments inputs input tensor initial_state list of initial state tensors to be passed to the first call of the cell mask binary tensor of shape batch_size timesteps indicating whether a given timestep should be masked an individual true entry indicates that the corresponding timestep should be utilized while a false entry indicates that the corresponding timestep should be ignored training python boolean indicating whether the layer should behave in training mode or in inference mode this argument is passed to the cell when calling it this is for use with cells that use dropout input shape 3 d tensor with shape batch_size timesteps features output shape if return_state a list of tensors the first tensor is the output the remaining tensors are the last states each with shape batch_size state_size where state_size could be a high dimension tensor shape if return_sequences 3d tensor with shape batch_size timesteps output_size masking this layer supports masking for input data with a variable number of timesteps to introduce masks to your data use a keras layers embedding layer with the mask_zero parameter set to true note on using statefulness in rnns you can set rnn layers to be stateful which means that the states computed for the samples in one batch will be reused as initial states for the samples in the next batch this assumes a one to one mapping between samples in different successive batches to enable statefulness specify stateful true in the layer constructor specify a fixed batch size for your model by passing if sequential model batch_input_shape to the first layer in your model else for functional model with 1 or more input layers batch_shape to all the first layers in your model this is the expected shape of your inputs including the batch size it should be a tuple of integers e g 32 10 100 specify shuffle false when calling fit to reset the states of your model call reset_states on either a specific layer or on your entire model note on specifying the initial state of rnns you can specify the initial state of rnn layers symbolically by calling them with the keyword argument initial_state the value of initial_state should be a tensor or list of tensors representing the initial state of the rnn layer you can specify the initial state of rnn layers numerically by calling reset_states with the keyword argument states the value of states should be a numpy array or list of numpy arrays representing the initial state of the rnn layer examples python from keras layers import rnn from keras import ops first let s define a rnn cell as a layer subclass class minimalrnncell keras layers layer def __init__ self units kwargs super __init__ kwargs self units units self state_size units def build self input_shape self kernel self add_weight shape input_shape 1 self units initializer uniform name kernel self recurrent_kernel self add_weight shape self units self units initializer uniform name recurrent_kernel self built true def call self inputs states prev_output states 0 h ops matmul inputs self kernel output h ops matmul prev_output self recurrent_kernel return output output let s use this cell in a rnn layer cell minimalrnncell 32 x keras input none 5 layer rnn cell y layer x here s how to use the cell to build a stacked rnn cells minimalrnncell 32 minimalrnncell 64 x keras input none 5 layer rnn cells y layer x if true the output for masked timestep will be zeros whereas in the false case output from previous timestep is returned for masked timestep time step masks must be the same for each input this is because the mask for an rnn is of size batch time_steps 1 and specifies which time steps should be skipped and a time step must be skipped for all inputs build cell if layer rnn expect the states in a list even if single state force the state to be a list in case it is a namedtuple eg lstmstatetuple compatibility alias rnn expect the states in a list even if single state cast states to compute dtype note that states may be deeply nested e g in the stacked cells case prepopulate the dropout state so that the inner_loop is stateless this is particularly important for jax backend
import tree

from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.layers.layer import Layer
from keras.layers.rnn.dropout_rnn_cell import DropoutRNNCell
from keras.layers.rnn.stacked_rnn_cells import StackedRNNCells
from keras.saving import serialization_lib
from keras.utils import tracking


@keras_export("keras.layers.RNN")
class RNN(Layer):
    

    def __init__(
        self,
        cell,
        return_sequences=False,
        return_state=False,
        go_backwards=False,
        stateful=False,
        unroll=False,
        zero_output_for_mask=False,
        **kwargs,
    ):
        if isinstance(cell, (list, tuple)):
            cell = StackedRNNCells(cell)
        if "call" not in dir(cell):
            raise ValueError(
                "Argument `cell` should have a `call` method. "
                f"Received: cell={cell}"
            )
        if "state_size" not in dir(cell):
            raise ValueError(
                "The RNN cell should have a `state_size` attribute "
                "(single integer or list of integers, "
                "one integer per RNN state). "
                f"Received: cell={cell}"
            )
        super().__init__(**kwargs)

        
        
        
        self.zero_output_for_mask = zero_output_for_mask
        self.cell = cell
        self.return_sequences = return_sequences
        self.return_state = return_state
        self.go_backwards = go_backwards
        self.stateful = stateful
        self.unroll = unroll

        self.supports_masking = True
        self.input_spec = None
        self.states = None

        state_size = getattr(self.cell, "state_size", None)
        if state_size is None:
            raise ValueError(
                "state_size must be specified as property on the RNN cell."
            )
        if not isinstance(state_size, (list, tuple, int)):
            raise ValueError(
                "state_size must be an integer, or a list/tuple of integers "
                "(one for each state tensor)."
            )
        if isinstance(state_size, int):
            self.state_size = [state_size]
            self.single_state = True
        else:
            self.state_size = list(state_size)
            self.single_state = False

    def compute_output_shape(self, sequences_shape, initial_state_shape=None):
        state_shape = [(sequences_shape[0], d) for d in self.state_size]
        output_size = getattr(self.cell, "output_size", None)
        if output_size is None:
            output_size = self.state_size[0]
        if not isinstance(output_size, int):
            raise ValueError("output_size must be an integer.")
        if self.return_sequences:
            output_shape = (sequences_shape[0], sequences_shape[1], output_size)
        else:
            output_shape = (sequences_shape[0], output_size)
        if self.return_state:
            return output_shape, *state_shape
        return output_shape

    def compute_mask(self, _, mask):
        
        
        
        
        mask = tree.flatten(mask)[0]
        output_mask = mask if self.return_sequences else None
        if self.return_state:
            state_mask = [None for _ in self.state_size]
            return [output_mask] + state_mask
        else:
            return output_mask

    def build(self, sequences_shape, initial_state_shape=None):
        
        step_input_shape = (sequences_shape[0],) + tuple(sequences_shape[2:])
        if isinstance(self.cell, Layer) and not self.cell.built:
            self.cell.build(step_input_shape)
            self.cell.built = True
        if self.stateful:
            if self.states is not None:
                self.reset_state()
            else:
                if sequences_shape[0] is None:
                    raise ValueError(
                        "When using `stateful=True` in a RNN, the "
                        "batch size must be static. Found dynamic "
                        f"batch size: sequence.shape={sequences_shape}"
                    )
                self._create_state_variables(sequences_shape[0])
        self.built = True

    @tracking.no_automatic_dependency_tracking
    def _create_state_variables(self, batch_size):
        with backend.name_scope(self.name, caller=self):
            self.states = tree.map_structure(
                lambda value: backend.Variable(
                    value,
                    trainable=False,
                    dtype=self.variable_dtype,
                    name="rnn_state",
                ),
                self.get_initial_state(batch_size),
            )

    def get_initial_state(self, batch_size):
        get_initial_state_fn = getattr(self.cell, "get_initial_state", None)
        if get_initial_state_fn:
            init_state = get_initial_state_fn(batch_size=batch_size)
        else:
            return [
                ops.zeros((batch_size, d), dtype=self.cell.compute_dtype)
                for d in self.state_size
            ]

        
        if not tree.is_nested(init_state):
            init_state = [init_state]
        
        
        return list(init_state)

    def reset_states(self):
        
        self.reset_state()

    def reset_state(self):
        if self.states is not None:
            for v in self.states:
                v.assign(ops.zeros_like(v))

    def inner_loop(self, sequences, initial_state, mask, training=False):
        cell_kwargs = {}
        if isinstance(self.cell, Layer) and self.cell._call_has_training_arg:
            cell_kwargs["training"] = training

        def step(inputs, states):
            output, new_states = self.cell(inputs, states, **cell_kwargs)
            if not tree.is_nested(new_states):
                new_states = [new_states]
            return output, new_states

        if not tree.is_nested(initial_state):
            initial_state = [initial_state]

        return backend.rnn(
            step,
            sequences,
            initial_state,
            go_backwards=self.go_backwards,
            mask=mask,
            unroll=self.unroll,
            input_length=sequences.shape[1],
            zero_output_for_mask=self.zero_output_for_mask,
            return_all_outputs=self.return_sequences,
        )

    def call(
        self,
        sequences,
        initial_state=None,
        mask=None,
        training=False,
    ):
        timesteps = sequences.shape[1]
        if self.unroll and timesteps is None:
            raise ValueError(
                "Cannot unroll a RNN if the "
                "time dimension is undefined. \n"
                "- If using a Sequential model, "
                "specify the time dimension by passing "
                "an `Input()` as your first layer.\n"
                "- If using the functional API, specify "
                "the time dimension by passing a `shape` "
                "or `batch_shape` argument to your `Input()`."
            )

        if initial_state is None:
            if self.stateful:
                initial_state = self.states
            else:
                initial_state = self.get_initial_state(
                    batch_size=ops.shape(sequences)[0]
                )
        
        if not tree.is_nested(initial_state):
            initial_state = [initial_state]
        initial_state = list(initial_state)

        
        
        
        initial_state = tree.map_structure(
            lambda x: backend.convert_to_tensor(
                x, dtype=self.cell.compute_dtype
            ),
            initial_state,
        )

        
        
        self._maybe_config_dropout_masks(self.cell, sequences, initial_state)

        last_output, outputs, states = self.inner_loop(
            sequences=sequences,
            initial_state=initial_state,
            mask=mask,
            training=training,
        )
        last_output = ops.cast(last_output, self.compute_dtype)
        outputs = ops.cast(outputs, self.compute_dtype)
        states = tree.map_structure(
            lambda x: ops.cast(x, dtype=self.compute_dtype), states
        )
        self._maybe_reset_dropout_masks(self.cell)

        if self.stateful:
            for self_state, state in zip(
                tree.flatten(self.states), tree.flatten(states)
            ):
                self_state.assign(state)

        if self.return_sequences:
            output = outputs
        else:
            output = last_output

        if self.return_state:
            if len(states) == 1:
                state = states[0]
                return output, state
            return output, *states
        return output

    def _maybe_config_dropout_masks(self, cell, input_sequence, input_state):
        step_input = input_sequence[:, 0, :]
        state = (
            input_state[0]
            if isinstance(input_state, (list, tuple))
            else input_state
        )
        if isinstance(cell, DropoutRNNCell):
            cell.get_dropout_mask(step_input)
            cell.get_recurrent_dropout_mask(state)
        if isinstance(cell, StackedRNNCells):
            for c, s in zip(cell.cells, input_state):
                self._maybe_config_dropout_masks(c, input_sequence, s)

    def _maybe_reset_dropout_masks(self, cell):
        if isinstance(cell, DropoutRNNCell):
            cell.reset_dropout_mask()
            cell.reset_recurrent_dropout_mask()
        if isinstance(cell, StackedRNNCells):
            for c in cell.cells:
                self._maybe_reset_dropout_masks(c)

    def get_config(self):
        config = {
            "return_sequences": self.return_sequences,
            "return_state": self.return_state,
            "go_backwards": self.go_backwards,
            "stateful": self.stateful,
            "unroll": self.unroll,
            "zero_output_for_mask": self.zero_output_for_mask,
        }
        config["cell"] = serialization_lib.serialize_keras_object(self.cell)
        base_config = super().get_config()
        return {**base_config, **config}

    @classmethod
    def from_config(cls, config, custom_objects=None):
        cell = serialization_lib.deserialize_keras_object(
            config.pop("cell"), custom_objects=custom_objects
        )
        layer = cls(cell, **config)
        return layer