class_docstr: Helper class to simplify implementing RNN models with TFModelV2.
  
  Instead of implementing forward(), you can implement forward_rnn() which
  takes batches with the time dimension added already.
  
  Here is an example implementation for a subclass
  ``MyRNNClass(RecurrentNetwork)``::
  
      def __init__(self, *args, **kwargs):
          super(MyModelClass, self).__init__(*args, **kwargs)
          cell_size = 256
  
          # Define input layers
          input_layer = tf.keras.layers.Input(
              shape=(None, obs_space.shape[0]))
          state_in_h = tf.keras.layers.Input(shape=(256, ))
          state_in_c = tf.keras.layers.Input(shape=(256, ))
          seq_in = tf.keras.layers.Input(shape=(), dtype=tf.int32)
  
          # Send to LSTM cell
          lstm_out, state_h, state_c = tf.keras.layers.LSTM(
              cell_size, return_sequences=True, return_state=True,
              name="lstm")(
                  inputs=input_layer,
                  mask=tf.sequence_mask(seq_in),
                  initial_state=[state_in_h, state_in_c])
          output_layer = tf.keras.layers.Dense(...)(lstm_out)
  
          # Create the RNN model
          self.rnn_model = tf.keras.Model(
              inputs=[input_layer, seq_in, state_in_h, state_in_c],
              outputs=[output_layer, state_h, state_c])
          self.rnn_model.summary()
class_docstr_tok: ["Helper","class","to","simplify","implementing","RNN","models","with","TFModelV2",".","Instead","of","implementing","forward","(",")",",","you","can","implement","forward_rnn","(",")","which","takes","batches","with","the","time","dimension","added","already",".","Here","is","an","example","implementation","for","a","subclass","``","MyRNNClass","(","RecurrentNetwork",")","``",":",":","def","__init__","(","self",",","*","args",",","*","*","kwargs",")",":","super","(","MyModelClass",",","self",")",".__init__","(","*","args",",","*","*","kwargs",")","cell_size","=","256","#","Define","input","layers","input_layer","=","tf.keras.layers.Input","(","shape=","(","None",",","obs_space.shape","[","0","]",")",")","state_in_h","=","tf.keras.layers.Input","(","shape=","(","256",",",")",")","state_in_c","=","tf.keras.layers.Input","(","shape=","(","256",",",")",")","seq_in","=","tf.keras.layers.Input","(","shape=","(",")",",","dtype=tf.int32",")","#","Send","to","LSTM","cell","lstm_out",",","state_h",",","state_c","=","tf.keras.layers.LSTM","(","cell_size",",","return_sequences=True",",","return_state=True",",","name=","''","lstm","''",")","(","inputs=input_layer",",","mask=tf.sequence_mask","(","seq_in",")",",","initial_state=","[","state_in_h",",","state_in_c","]",")","output_layer","=","tf.keras.layers.Dense","(","...",")","(","lstm_out",")","#","Create","the","RNN","model","self.rnn_model","=","tf.keras.Model","(","inputs=","[","input_layer",",","seq_in",",","state_in_h",",","state_in_c","]",",","outputs=","[","output_layer",",","state_h",",","state_c","]",")","self.rnn_model.summary","(",")"]
docstr: Adds time dimension to batch before sending inputs to forward_rnn().
  
  You should implement forward_rnn() in your subclass.
docstr_tok: ["Adds","time","dimension","to","batch","before","sending","inputs","to","forward_rnn","(",")",".","You","should","implement","forward_rnn","(",")","in","your","subclass","."]
code: def forward(
      self,
      input_dict: Dict[str, TensorType],
      state: List[TensorType],
      seq_lens: TensorType,
  ) -> Tuple[TensorType, List[TensorType]]:
      """Adds time dimension to batch before sending inputs to forward_rnn().
  
      You should implement forward_rnn() in your subclass."""
      # Creating a __init__ function that acts as a passthrough and adding the warning
      # there led to errors probably due to the multiple inheritance. We encountered
      # the same error if we add the Deprecated decorator. We therefore add the
      # deprecation warning here.
      if log_once("recurrent_network_tf"):
          deprecation_warning(
              old="ray.rllib.models.tf.recurrent_net.RecurrentNetwork"
          )
      assert seq_lens is not None
      flat_inputs = input_dict["obs_flat"]
      inputs = add_time_dimension(
          padded_inputs=flat_inputs, seq_lens=seq_lens, framework="tf"
      )
      output, new_state = self.forward_rnn(
          inputs,
          state,
          seq_lens,
      )
      return tf.reshape(output, [-1, self.num_outputs]), new_state
  
imports: import numpy
  import gymnasium
  from gymnasium.spaces import Discrete, MultiDiscrete
  import logging
  import tree
  from typing import Dict, List, Tuple
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.policy.rnn_sequencing import add_time_dimension
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils.annotations import override, DeveloperAPI
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space
  from ray.rllib.utils.tf_utils import flatten_inputs_to_1d_tensor, one_hot
  from ray.rllib.utils.typing import ModelConfigDict, TensorType
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.util.debug import log_once
label_desc: Use image node_id 1 for calling the RecurrentNetwork obj's underlying member method code with example usage: obj.forward(input_dict, state, seq_lens) and returns: tf,  new_state