(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)
return h.view(length, batch, -1), torch.stack(c_final, dim=1).view(batch, -1)
Elementwise forward operation of SRU in pure Python.
Elementwise forward operation of SRU in pure Python.
Elementwise forward operation of SRU in pure Python .
def elementwise_recurrence_naive(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]:
    """Elementwise forward operation of SRU in pure Python.

    """
    if torch.is_grad_enabled():
        warnings.warn("Running SRU on CPU with grad_enabled=True. Are you sure?")
    else:
        return elementwise_recurrence_inference(U, x, weight_c, bias, c_init,
                                                activation_type, hidden_size,
                                                bidirectional, has_skip_term,
                                                scale_x, dropout_mask_c, mask_pad)

    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
    d = hidden_size
    is_custom = weight_c.dim() > 1

    U = U.contiguous().view(length, batch, bidir, d, k)

    if is_custom:
        weight_c = weight_c.view(length, batch, bidir, d, 2)
        forget_wc = weight_c[..., 0]
        reset_wc = weight_c[..., 1]
    else:
        forget_wc, reset_wc = weight_c.view(2, bidir, d)

    forget_bias, reset_bias = bias.view(2, bidir, d)

    if not has_skip_term:
        x_prime = None
    elif k == 3:
        x_prime = x.view(length, batch, bidir, d)
        x_prime = x_prime * scale_x if scale_x is not None else x_prime
    else:
        x_prime = U[..., 3]

    if c_init is None:
        c_init = x.new_zeros(size=(batch, bidir, d))
    else:
        c_init = c_init.view(batch, bidir, d)

    mask_pad = mask_pad.view(length, batch, 1).float() if mask_pad is not None else None
    mask_c = dropout_mask_c.view(batch, bidir, d) if dropout_mask_c is not None else None

    h = x.new_zeros(length, batch, bidir, d)
    c_final = []
    for di in range(bidir):
        time_seq = range(length) if di == 0 else range(length - 1, -1, -1)
        mask_c_ = 1 if mask_c is None else mask_c[:, di, :]
        c_prev = c_init[:, di, :]
        fb, rb = forget_bias[di], reset_bias[di]
        if is_custom:
            fw = forget_wc[:, :, di, :].chunk(length)
            rw = reset_wc[:, :, di, :].chunk(length)
        else:
            fw = forget_wc[di].expand(batch, d)  # type: ignore
            rw = reset_wc[di].expand(batch, d)  # type: ignore
        u0 = U[:, :, di, :, 0].chunk(length)
        u1 = (U[:, :, di, :, 1] + fb).chunk(length)
        u2 = (U[:, :, di, :, 2] + rb).chunk(length)
        if x_prime is not None:
            xp = x_prime[:, :, di, :].chunk(length)

        for t in time_seq:
            if is_custom:
                forget_t = (u1[t] + c_prev*fw[t]).sigmoid()
                reset_t = (u2[t] + c_prev*rw[t]).sigmoid()
            else:
                forget_t = (u1[t] + c_prev*fw).sigmoid()
                reset_t = (u2[t] + c_prev*rw).sigmoid()
            c_t = u0[t] + (c_prev - u0[t]) * forget_t
            if mask_pad is not None:
                c_t = c_t * (1-mask_pad[t]) + c_prev * mask_pad[t]
            c_prev = c_t

            if activation_type == 0:
                g_c_t = c_t
            elif activation_type == 1:
                g_c_t = c_t.tanh()
            else:
                raise ValueError('Activation type must be 0 or 1, not {}'.format(activation_type))

            if x_prime is not None:
                h_t = xp[t] + (g_c_t - xp[t]) * mask_c_ * reset_t
            else:
                h_t = g_c_t * mask_c_ * reset_t
            if mask_pad is not None:
                h_t = h_t * (1-mask_pad[t])
            h[t, :, di, :] = h_t
        c_final.append(c_t.view(batch, d))

    return h.view(length, batch, -1), torch.stack(c_final, dim=1).view(batch, -1)
def elementwise_recurrence_naive ( 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 ] : if torch . is_grad_enabled ( ) : warnings . warn ( "Running SRU on CPU with grad_enabled=True. Are you sure?" ) else : return elementwise_recurrence_inference ( U , x , weight_c , bias , c_init , activation_type , hidden_size , bidirectional , has_skip_term , scale_x , dropout_mask_c , mask_pad ) 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 d = hidden_size is_custom = weight_c . dim ( ) > 1 U = U . contiguous ( ) . view ( length , batch , bidir , d , k ) if is_custom : weight_c = weight_c . view ( length , batch , bidir , d , 2 ) forget_wc = weight_c [ ... , 0 ] reset_wc = weight_c [ ... , 1 ] else : forget_wc , reset_wc = weight_c . view ( 2 , bidir , d ) forget_bias , reset_bias = bias . view ( 2 , bidir , d ) if not has_skip_term : x_prime = None elif k == 3 : x_prime = x . view ( length , batch , bidir , d ) x_prime = x_prime * scale_x if scale_x is not None else x_prime else : x_prime = U [ ... , 3 ] if c_init is None : c_init = x . new_zeros ( size = ( batch , bidir , d ) ) else : c_init = c_init . view ( batch , bidir , d ) mask_pad = mask_pad . view ( length , batch , 1 ) . float ( ) if mask_pad is not None else None mask_c = dropout_mask_c . view ( batch , bidir , d ) if dropout_mask_c is not None else None h = x . new_zeros ( length , batch , bidir , d ) c_final = [ ] for di in range ( bidir ) : time_seq = range ( length ) if di == 0 else range ( length - 1 , - 1 , - 1 ) mask_c_ = 1 if mask_c is None else mask_c [ : , di , : ] c_prev = c_init [ : , di , : ] fb , rb = forget_bias [ di ] , reset_bias [ di ] if is_custom : fw = forget_wc [ : , : , di , : ] . chunk ( length ) rw = reset_wc [ : , : , di , : ] . chunk ( length ) else : fw = forget_wc [ di ] . expand ( batch , d ) # type: ignore rw = reset_wc [ di ] . expand ( batch , d ) # type: ignore u0 = U [ : , : , di , : , 0 ] . chunk ( length ) u1 = ( U [ : , : , di , : , 1 ] + fb ) . chunk ( length ) u2 = ( U [ : , : , di , : , 2 ] + rb ) . chunk ( length ) if x_prime is not None : xp = x_prime [ : , : , di , : ] . chunk ( length ) for t in time_seq : if is_custom : forget_t = ( u1 [ t ] + c_prev * fw [ t ] ) . sigmoid ( ) reset_t = ( u2 [ t ] + c_prev * rw [ t ] ) . sigmoid ( ) else : forget_t = ( u1 [ t ] + c_prev * fw ) . sigmoid ( ) reset_t = ( u2 [ t ] + c_prev * rw ) . sigmoid ( ) c_t = u0 [ t ] + ( c_prev - u0 [ t ] ) * forget_t if mask_pad is not None : c_t = c_t * ( 1 - mask_pad [ t ] ) + c_prev * mask_pad [ t ] c_prev = c_t if activation_type == 0 : g_c_t = c_t elif activation_type == 1 : g_c_t = c_t . tanh ( ) else : raise ValueError ( 'Activation type must be 0 or 1, not {}' . format ( activation_type ) ) if x_prime is not None : h_t = xp [ t ] + ( g_c_t - xp [ t ] ) * mask_c_ * reset_t else : h_t = g_c_t * mask_c_ * reset_t if mask_pad is not None : h_t = h_t * ( 1 - mask_pad [ t ] ) h [ t , : , di , : ] = h_t c_final . append ( c_t . view ( batch , d ) ) return h . view ( length , batch , - 1 ) , torch . stack ( c_final , dim = 1 ) . view ( batch , - 1 )