class_docstr: Wraps a Keras RNN/LSTM/GRU layer to make network state more consistent.
docstr: Perform the computation.
  
  Args:
    inputs: A tuple containing tensors in batch-major format, each shaped
      `[batch_size, n, ...]`.
    initial_state: (Optional) An initial state for the wrapped layer. If not
      provided, `get_initial_state()` is used instead.
    mask: The mask to pass down to the wrapped layer.
    training: Whether the output is being used for training.
  
  Returns:
    A 2-tuple `(outputs, final_state)` where:
  
     - `outputs` contains the outputs for all time steps of the unroll; this
       is typically a tensor shaped `[batch_size, n, ...]`.
     - `final_state` contains the final state.
docstr_tok: ["Perform","the","computation",".","Args",":","inputs",":","A","tuple","containing","tensors","in","batch-major","format",",","each","shaped","`","[","batch_size",",","n",",","...","]","`",".","initial_state",":","(","Optional",")","An","initial","state","for","the","wrapped","layer",".","If","not","provided",",","`","get_initial_state","(",")","`","is","used","instead",".","mask",":","The","mask","to","pass","down","to","the","wrapped","layer",".","training",":","Whether","the","output","is","being","used","for","training",".","Returns",":","A","2-tuple","`","(","outputs",",","final_state",")","`","where",":","-","`","outputs","`","contains","the","outputs","for","all","time","steps","of","the","unroll",";","this","is","typically","a","tensor","shaped","`","[","batch_size",",","n",",","...","]","`",".","-","`","final_state","`","contains","the","final","state","."]
code: def call(self, inputs, initial_state=None, mask=None, training=False):
      """Perform the computation.
  
      Args:
        inputs: A tuple containing tensors in batch-major format, each shaped
          `[batch_size, n, ...]`.
        initial_state: (Optional) An initial state for the wrapped layer. If not
          provided, `get_initial_state()` is used instead.
        mask: The mask to pass down to the wrapped layer.
        training: Whether the output is being used for training.
  
      Returns:
        A 2-tuple `(outputs, final_state)` where:
  
         - `outputs` contains the outputs for all time steps of the unroll; this
           is typically a tensor shaped `[batch_size, n, ...]`.
         - `final_state` contains the final state.
      """
      inputs_flat = [
          tf.convert_to_tensor(x, name="input", dtype_hint=self.dtype)
          for x in tf.nest.flatten(inputs)
      ]
      has_time_axis = all(
          [
              x.shape.ndims is None or x.shape.ndims > 2
              for x in inputs_flat
          ]
      )
      if not has_time_axis:
          inputs_flat = [tf.expand_dims(t, axis=1) for t in inputs_flat]
      inputs = tf.nest.pack_sequence_as(inputs, inputs_flat)
  
      if not common.safe_has_state(initial_state):
          initial_state = self._layer.get_initial_state(inputs)
  
      outputs = self._layer(
          inputs,
          initial_state=initial_state,
          mask=mask,
          training=training,
      )
  
      output, new_state = outputs[0], outputs[1:]
  
      # Keras RNN's outputs[1:] does not match the nest structure of its cells'
      # state_size property.  Restructure the output state to match.
      new_state = tf.nest.pack_sequence_as(
          self.state_size, tf.nest.flatten(new_state)
      )
  
      if not has_time_axis:
          output = tf.nest.map_structure(
              lambda t: tf.squeeze(t, axis=1), output
          )
  
      # Outputs are in output, and state is in outputs[1:]
      return output, new_state
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.utils import common
label_desc: Use image node_id 15 for calling the RNNWrapper obj's underlying member method code with example usage: obj.call(inputs, initial_state, mask, training) and returns: output,  new_state