class_docstr: Base class for all RNN-like sequence models.
  
  DP-friendly drop-in replacement of the ``torch.nn.RNNBase`` module.
  After training this module can be exported and loaded by the original ``torch.nn``
  implementation for inference.
  
  This module implements multi-layer (Type-2, see
  [this issue](https://github.com/pytorch/pytorch/issues/4930#issuecomment-361851298))
  bi-directional sequential model based on abstract cell.
  Cell should be a subclass of ``DPRNNCellBase``.
  
  Limitations:
  - proj_size > 0 is not implemented
  - this implementation doesn't use cuDNN
class_docstr_tok: ["Base","class","for","all","RNN-like","sequence","models",".","DP-friendly","drop-in","replacement","of","the","``","torch.nn.RNNBase","``","module",".","After","training","this","module","can","be","exported","and","loaded","by","the","original","``","torch.nn","``","implementation","for","inference",".","This","module","implements","multi-layer","(","Type-2",",","see","[","this","issue","]","(","https",":","\/\/github.com\/pytorch\/pytorch\/issues\/4930","#","issuecomment-361851298",")",")","bi-directional","sequential","model","based","on","abstract","cell",".","Cell","should","be","a","subclass","of","``","DPRNNCellBase","``",".","Limitations",":","-","proj_size",">","0","is","not","implemented","-","this","implementation","does","n't","use","cuDNN"]
docstr: Forward pass of a full RNN, containing one or many single- or bi-directional layers.
  Implemented for an abstract cell type.
  
  Note: ``proj_size > 0`` is not supported here.
  Cell state size is always equal to hidden state size.
  
  Inputs: input, h_0/(h_0, c_0)
      input: Input sequence. Tensor of shape ``[T, B, D]`` (``[B, T, D]`` if ``batch_first=True``)
             or PackedSequence.
      h_0: Initial hidden state for each element in the batch. Tensor of shape ``[L*P, B, H]``. Default to zeros.
      c_0: Initial cell state for each element in the batch. Only for cell types with an additional state.
           Tensor of shape ``[L*P, B, H]``. Default to zeros.
  
  Outputs: output, h_n/(h_n, c_n)
      output: Output features (``h_t``) from the last layer of the model for each ``t``. Tensor of
          shape ``[T, B, P*H]`` (``[B, T, P*H]`` if ``batch_first=True``), or PackedSequence.
      h_n: Final hidden state for each element in the batch. Tensor of shape ``[L*P, B, H]``.
      c_n: Final cell state for each element in the batch. Tensor of shape ``[L*P, B, H]``.
  
  where
      T = sequence length
      B = batch size
      D = input_size
      H = hidden_size
      L = num_layers
      P = num_directions (2 if `bidirectional=True` else 1)
docstr_tok: ["Forward","pass","of","a","full","RNN",",","containing","one","or","many","single-","or","bi-directional","layers",".","Implemented","for","an","abstract","cell","type",".","Note",":","``","proj_size",">","0","``","is","not","supported","here",".","Cell","state","size","is","always","equal","to","hidden","state","size",".","Inputs",":","input",",","h_0\/","(","h_0",",","c_0",")","input",":","Input","sequence",".","Tensor","of","shape","``","[","T",",","B",",","D","]","``","(","``","[","B",",","T",",","D","]","``","if","``","batch_first=True","``",")","or","PackedSequence",".","h_0",":","Initial","hidden","state","for","each","element","in","the","batch",".","Tensor","of","shape","``","[","L","*","P",",","B",",","H","]","``",".","Default","to","zeros",".","c_0",":","Initial","cell","state","for","each","element","in","the","batch",".","Only","for","cell","types","with","an","additional","state",".","Tensor","of","shape","``","[","L","*","P",",","B",",","H","]","``",".","Default","to","zeros",".","Outputs",":","output",",","h_n\/","(","h_n",",","c_n",")","output",":","Output","features","(","``","h_t","``",")","from","the","last","layer","of","the","model","for","each","``","t","``",".","Tensor","of","shape","``","[","T",",","B",",","P","*","H","]","``","(","``","[","B",",","T",",","P","*","H","]","``","if","``","batch_first=True","``",")",",","or","PackedSequence",".","h_n",":","Final","hidden","state","for","each","element","in","the","batch",".","Tensor","of","shape","``","[","L","*","P",",","B",",","H","]","``",".","c_n",":","Final","cell","state","for","each","element","in","the","batch",".","Tensor","of","shape","``","[","L","*","P",",","B",",","H","]","``",".","where","T","=","sequence","length","B","=","batch","size","D","=","input_size","H","=","hidden_size","L","=","num_layers","P","=","num_directions","(","2","if","`","bidirectional=True","`","else","1",")"]
code: def forward(
      self,
      input: Union[Tensor, PackedSequence],
      state_init: Optional[Union[Tensor, Tuple[Tensor, Tensor]]] = None,
  ) -> Tuple[
      Union[Tensor, PackedSequence],
      Union[Tensor, Tuple[Tensor, Tensor]],
  ]:
      """
      Forward pass of a full RNN, containing one or many single- or bi-directional layers.
      Implemented for an abstract cell type.
  
      Note: ``proj_size > 0`` is not supported here.
      Cell state size is always equal to hidden state size.
  
      Inputs: input, h_0/(h_0, c_0)
          input: Input sequence. Tensor of shape ``[T, B, D]`` (``[B, T, D]`` if ``batch_first=True``)
                 or PackedSequence.
          h_0: Initial hidden state for each element in the batch. Tensor of shape ``[L*P, B, H]``. Default to zeros.
          c_0: Initial cell state for each element in the batch. Only for cell types with an additional state.
               Tensor of shape ``[L*P, B, H]``. Default to zeros.
  
      Outputs: output, h_n/(h_n, c_n)
          output: Output features (``h_t``) from the last layer of the model for each ``t``. Tensor of
              shape ``[T, B, P*H]`` (``[B, T, P*H]`` if ``batch_first=True``), or PackedSequence.
          h_n: Final hidden state for each element in the batch. Tensor of shape ``[L*P, B, H]``.
          c_n: Final cell state for each element in the batch. Tensor of shape ``[L*P, B, H]``.
  
      where
          T = sequence length
          B = batch size
          D = input_size
          H = hidden_size
          L = num_layers
          P = num_directions (2 if `bidirectional=True` else 1)
      """
      num_directions = 2 if self.bidirectional else 1
  
      is_packed = isinstance(input, PackedSequence)
      if is_packed:
          (
              input_data,
              batch_sizes,
              sorted_indices,
              unsorted_indices,
          ) = input
          dtype, device = input_data.dtype, input_data.device
  
          x = input_data.split(tuple(batch_sizes))  # tuple T x [B, D]
  
          seq_length = len(batch_sizes)
          max_batch_size = int(batch_sizes[0])
  
          for cell in self.cells:
              cell.set_max_batch_length(max_batch_size)
      else:
          dtype, device = input.dtype, input.device
          batch_sizes = None
          sorted_indices = None
          unsorted_indices = None
  
          # Rearrange batch dim. Batch is by default in second dimension.
          if self.batch_first:
              input = input.transpose(0, 1)
  
          x = input  # [T, B, D]
  
          seq_length = x.shape[0]
          max_batch_size = x.shape[1]
  
      if self.has_cell_state:
          h_0s, c_0s = state_init or (None, None)
      else:
          h_0s, c_0s = state_init, None
  
      if h_0s is None:
          h_0s = torch.zeros(  # [L*P, B, H]
              self.num_layers * num_directions,
              max_batch_size,
              self.hidden_size,
              dtype=dtype,
              device=device,
          )
      else:
          h_0s = apply_permutation(h_0s, 1, sorted_indices)
  
      if self.has_cell_state:
          if c_0s is None:
              c_0s = torch.zeros(  # [L*P, B, H]
                  self.num_layers * num_directions,
                  max_batch_size,
                  self.hidden_size,
                  dtype=dtype,
                  device=device,
              )
          else:
              c_0s = apply_permutation(c_0s, 1, sorted_indices)
      else:
          c_0s = [None] * len(h_0s)
  
      hs = []
      cs = []  # list of None if no cell state
      output = None
  
      for layer, directions in self.iterate_layers(
          self.cells, h_0s, c_0s
      ):
          layer_outs = []
  
          for direction, (cell, h0, c0) in directions:
              # apply single direction layer (with dropout)
              out_layer, h, c = self.forward_layer(
                  x
                  if layer == 0
                  else output,  # [T, B, D/H/2H] / tuple T x [B, D/H/2H]
                  h0,  # [B, H]
                  c0,
                  batch_sizes,
                  cell=cell,
                  max_batch_size=max_batch_size,
                  seq_length=seq_length,
                  is_packed=is_packed,
                  reverse_layer=(direction == 1),
              )
  
              hs.append(h)  # h: [B, H]
              cs.append(c)
              layer_outs.append(
                  out_layer
              )  # out_layer: [T, B, H] / tuple T x [B, H]
  
          if is_packed:
              output = [  # tuple T x [B, P*H]
                  torch.cat(
                      [layer_out[i] for layer_out in layer_outs], dim=1
                  )
                  for i in range(seq_length)
              ]
          else:
              output = torch.cat(layer_outs, dim=2)  # [T, B, P*H]
  
      if is_packed:
          packed_data = torch.cat(output, dim=0)  # [TB, P*H]
          output = PackedSequence(
              packed_data, batch_sizes, sorted_indices, unsorted_indices
          )
      else:
          # Rearrange batch dim back
          if self.batch_first:
              output = output.transpose(0, 1)
  
      hs = torch.stack(hs, dim=0).to(device)  # [L*P, B, H]
      hs = apply_permutation(hs, 1, unsorted_indices)
      if self.has_cell_state:
          cs = torch.stack(cs, dim=0).to(device)  # [L*P, B, H]
          cs = apply_permutation(cs, 1, unsorted_indices)
  
      hidden = (hs, cs) if self.has_cell_state else hs
  
      return output, hidden
  
imports: import math
  import numbers
  import warnings
  from typing import List, Optional, Tuple, Type, Union
  import torch
  import torch.nn
  from torch import Tensor
  from torch.nn.utils.rnn import PackedSequence
  from ..utils.packed_sequences import compute_seq_lengths
  from .param_rename import RenameParamsMixin
label_desc: Use image node_id 2 for calling the DPRNNBase obj's underlying member method code with example usage: obj.forward(input, state_init) and returns: output,  hidden