(self, input: Tensor,
                c0: Optional[Tensor] = None,
                mask_pad: Optional[Tensor] = None)
The forward method of SRU module

        Parameters
        ----------
        input: Tensor
            the input feature. shape: (length, batch_size, input_size)
        c0: Tensor, optional
            the initial internal hidden state. shape: (num_layers,
            batch_size, output_size) where
            output_size = hidden_size * num_direction
        mask_pad: Tensor, optional
            the mask where a non-zero value indicates if an input token
            is pad token that should be ignored in forward and backward
            computation. shape: (length, batch_size)

        Returns
        ----------
        h: Tensor
            the output hidden state. shape: (length, batch_size,
            output_size) where
            output_size = hidden_size * num_direction
        c: Tensor
            the last internal hidden state. shape: (num_layers,
            batch_size, output_size), or (num_layers * num_directions,
            batch_size, hidden_size) if `nn_rnn_compatible_return` is
            set `True`
The forward method of SRU module
The forward method of SRU module
def forward(self, input: Tensor,
                c0: Optional[Tensor] = None,
                mask_pad: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
        """The forward method of SRU module

        Parameters
        ----------
        input: Tensor
            the input feature. shape: (length, batch_size, input_size)
        c0: Tensor, optional
            the initial internal hidden state. shape: (num_layers,
            batch_size, output_size) where
            output_size = hidden_size * num_direction
        mask_pad: Tensor, optional
            the mask where a non-zero value indicates if an input token
            is pad token that should be ignored in forward and backward
            computation. shape: (length, batch_size)

        Returns
        ----------
        h: Tensor
            the output hidden state. shape: (length, batch_size,
            output_size) where
            output_size = hidden_size * num_direction
        c: Tensor
            the last internal hidden state. shape: (num_layers,
            batch_size, output_size), or (num_layers * num_directions,
            batch_size, hidden_size) if `nn_rnn_compatible_return` is
            set `True`

        """
        # unpack packed, if input is packed. packing and then unpacking will be slower than not
        # packing at all, but makes SRU usage compatible with nn.RNN usage
        orig_input = input
        if isinstance(orig_input, PackedSequence):
            input, lengths = nn.utils.rnn.pad_packed_sequence(input)
            max_length = lengths.max().item()
            mask_pad = torch.ByteTensor([[0] * length + [1] * (max_length - length)
                                        for length in lengths.tolist()])
            mask_pad = mask_pad.to(input.device).transpose(0, 1).contiguous()

        # The dimensions of `input` should be: `(sequence_length, batch_size, input_size)`.
        if input.dim() != 3:
            raise ValueError("There must be 3 dimensions for (length, batch_size, input_size)")

        if c0 is None:
            zeros = torch.zeros(input.size(1), self.output_size, dtype=input.dtype,
                                device=input.device)
            c0_ = [zeros for i in range(self.num_layers)]
        else:
            # The dimensions of `c0` should be: `(num_layers, batch_size, hidden_size * dir_)`.
            if c0.dim() != 3:
                raise ValueError("c0 must be 3 dim (num_layers, batch_size, output_size)")
            c0_ = [x.squeeze(0) for x in c0.chunk(self.num_layers, 0)]

        if self.input_to_hidden is None:
            prevx = input
        else:
            prevx = self.input_to_hidden(input)
        lstc = []
        i = 0
        for rnn in self.rnn_lst:
            h, c = rnn(prevx, c0_[i], mask_pad=mask_pad)
            prevx = h
            lstc.append(c)
            i += 1

        lstc_stack = torch.stack(lstc)
        if self.nn_rnn_compatible_return:
            batch_size = input.size(1)
            lstc_stack = lstc_stack.view(self.num_layers, batch_size,
                                         self.num_directions, self.hidden_size)
            lstc_stack = lstc_stack.transpose(1, 2).contiguous()
            lstc_stack = lstc_stack.view(self.num_layers * self.num_directions,
                                         batch_size, self.hidden_size)

        if isinstance(orig_input, PackedSequence):
            prevx = nn.utils.rnn.pack_padded_sequence(prevx, lengths, enforce_sorted=False)
            return prevx, lstc_stack
        else:
            return prevx, lstc_stack
def forward ( self , input : Tensor , c0 : Optional [ Tensor ] = None , mask_pad : Optional [ Tensor ] = None ) -> Tuple [ Tensor , Tensor ] : # unpack packed, if input is packed. packing and then unpacking will be slower than not # packing at all, but makes SRU usage compatible with nn.RNN usage orig_input = input if isinstance ( orig_input , PackedSequence ) : input , lengths = nn . utils . rnn . pad_packed_sequence ( input ) max_length = lengths . max ( ) . item ( ) mask_pad = torch . ByteTensor ( [ [ 0 ] * length + [ 1 ] * ( max_length - length ) for length in lengths . tolist ( ) ] ) mask_pad = mask_pad . to ( input . device ) . transpose ( 0 , 1 ) . contiguous ( ) # The dimensions of `input` should be: `(sequence_length, batch_size, input_size)`. if input . dim ( ) != 3 : raise ValueError ( "There must be 3 dimensions for (length, batch_size, input_size)" ) if c0 is None : zeros = torch . zeros ( input . size ( 1 ) , self . output_size , dtype = input . dtype , device = input . device ) c0_ = [ zeros for i in range ( self . num_layers ) ] else : # The dimensions of `c0` should be: `(num_layers, batch_size, hidden_size * dir_)`. if c0 . dim ( ) != 3 : raise ValueError ( "c0 must be 3 dim (num_layers, batch_size, output_size)" ) c0_ = [ x . squeeze ( 0 ) for x in c0 . chunk ( self . num_layers , 0 ) ] if self . input_to_hidden is None : prevx = input else : prevx = self . input_to_hidden ( input ) lstc = [ ] i = 0 for rnn in self . rnn_lst : h , c = rnn ( prevx , c0_ [ i ] , mask_pad = mask_pad ) prevx = h lstc . append ( c ) i += 1 lstc_stack = torch . stack ( lstc ) if self . nn_rnn_compatible_return : batch_size = input . size ( 1 ) lstc_stack = lstc_stack . view ( self . num_layers , batch_size , self . num_directions , self . hidden_size ) lstc_stack = lstc_stack . transpose ( 1 , 2 ) . contiguous ( ) lstc_stack = lstc_stack . view ( self . num_layers * self . num_directions , batch_size , self . hidden_size ) if isinstance ( orig_input , PackedSequence ) : prevx = nn . utils . rnn . pack_padded_sequence ( prevx , lengths , enforce_sorted = False ) return prevx , lstc_stack else : return prevx , lstc_stack