(U: Tensor,
                                     x: Tensor,
                                     weight_c: Tensor,
                                     bias: Tensor,
                                     c_init: Tensor,
                                     activation_type: int,
                                     hidden_size: int,
                                     bidirectional: bool,
                                     has_skip_term: bool,
                                     scale_x: Optional[Tensor] = None,
                                     dropout_mask_c: Optional[Tensor] = None,
                                     mask_pad: Optional[Tensor] = None)
Torchscripted elementwise forward operation of SRU for inference.
Torchscripted elementwise forward operation of SRU for inference.
Torchscripted elementwise forward operation of SRU for inference .
def elementwise_recurrence_inference(U: Tensor,
                                     x: Tensor,
                                     weight_c: Tensor,
                                     bias: Tensor,
                                     c_init: Tensor,
                                     activation_type: int,
                                     hidden_size: int,
                                     bidirectional: bool,
                                     has_skip_term: bool,
                                     scale_x: Optional[Tensor] = None,
                                     dropout_mask_c: Optional[Tensor] = None,
                                     mask_pad: Optional[Tensor] = None) -> List[Tensor]:
    """Torchscripted elementwise forward operation of SRU for inference.

    """
    assert dropout_mask_c is None, "Dropout mask cannot be set during inference"
    bidir = 2 if bidirectional else 1
    length = x.size(0) if x.dim() == 3 else 1
    batch = x.size(-2)
    k = U.size(-1) // hidden_size // bidir
    is_custom = weight_c.dim() > 1
    mask_pad = None if mask_pad is None else mask_pad.to(dtype=torch.bool).contiguous()
    if U.is_cuda:
        h, last_hidden, c = elementwise_recurrence_forward(
            U,
            x,
            weight_c,
            bias,
            c_init,
            activation_type,
            hidden_size,
            bidirectional,
            has_skip_term,
            scale_x,
            dropout_mask_c,
            mask_pad
        )
        return h, last_hidden
    elif not bidirectional:
        return torch.ops.sru_cpu.cpu_forward(
            U.contiguous(),
            x.contiguous(),
            weight_c.contiguous(),
            bias.contiguous(),
            c_init.contiguous(),
            mask_pad,
            length,
            batch,
            hidden_size,
            k,
            activation_type,
            has_skip_term,
            scale_x.item() if scale_x is not None else 1.0,
            is_custom
        )
    else:
        return torch.ops.sru_cpu.cpu_bi_forward(
            U.contiguous(),
            x.contiguous(),
            weight_c.contiguous(),
            bias.contiguous(),
            c_init.contiguous(),
            mask_pad,
            length,
            batch,
            hidden_size,
            k,
            activation_type,
            has_skip_term,
            scale_x.item() if scale_x is not None else 1.0,
            is_custom
        )
def elementwise_recurrence_inference ( U : Tensor , x : Tensor , weight_c : Tensor , bias : Tensor , c_init : Tensor , activation_type : int , hidden_size : int , bidirectional : bool , has_skip_term : bool , scale_x : Optional [ Tensor ] = None , dropout_mask_c : Optional [ Tensor ] = None , mask_pad : Optional [ Tensor ] = None ) -> List [ Tensor ] : assert dropout_mask_c is None , "Dropout mask cannot be set during inference" bidir = 2 if bidirectional else 1 length = x . size ( 0 ) if x . dim ( ) == 3 else 1 batch = x . size ( - 2 ) k = U . size ( - 1 ) // hidden_size // bidir is_custom = weight_c . dim ( ) > 1 mask_pad = None if mask_pad is None else mask_pad . to ( dtype = torch . bool ) . contiguous ( ) if U . is_cuda : h , last_hidden , c = elementwise_recurrence_forward ( U , x , weight_c , bias , c_init , activation_type , hidden_size , bidirectional , has_skip_term , scale_x , dropout_mask_c , mask_pad ) return h , last_hidden elif not bidirectional : return torch . ops . sru_cpu . cpu_forward ( U . contiguous ( ) , x . contiguous ( ) , weight_c . contiguous ( ) , bias . contiguous ( ) , c_init . contiguous ( ) , mask_pad , length , batch , hidden_size , k , activation_type , has_skip_term , scale_x . item ( ) if scale_x is not None else 1.0 , is_custom ) else : return torch . ops . sru_cpu . cpu_bi_forward ( U . contiguous ( ) , x . contiguous ( ) , weight_c . contiguous ( ) , bias . contiguous ( ) , c_init . contiguous ( ) , mask_pad , length , batch , hidden_size , k , activation_type , has_skip_term , scale_x . item ( ) if scale_x is not None else 1.0 , is_custom )