pub fn forward_t(
        &self,
        hidden_states: &Tensor,
        attention_mask: Option<&LogitsMask>,
        train: bool,
    ) -> (Tensor, Tensor) {
        let mixed_key_layer = self.key.forward(hidden_states);
        let mixed_query_layer = self.query.forward(hidden_states);
        let mixed_value_layer = self.value.forward(hidden_states);

        let query_layer = self.transpose_for_scores(&mixed_query_layer);
        let key_layer = self.transpose_for_scores(&mixed_key_layer);
        let value_layer = self.transpose_for_scores(&mixed_value_layer);

        // Get the raw attention scores.
        let mut attention_scores = query_layer.matmul(&key_layer.transpose(-1, -2));
        attention_scores /= (self.attention_head_size as f64).sqrt();

        if let Some(mask) = attention_mask {
            attention_scores += &**mask;
        }

        // Convert the raw attention scores into a probability distribution.
        let attention_probs = attention_scores.softmax(-1, Kind::Float);

        // Drop out entire tokens to attend to, following the original
        // transformer paper.
        let attention_probs = self.dropout.forward_t(&attention_probs, train);

        let context_layer = attention_probs.matmul(&value_layer);

        let context_layer = context_layer.permute(&[0, 2, 1, 3]).contiguous();
        let mut new_context_layer_shape = context_layer.size();
        new_context_layer_shape.splice(
            new_context_layer_shape.len() - 2..,
            iter::once(self.all_head_size),
        );
        let context_layer = context_layer.view_(&new_context_layer_shape);

        (context_layer, attention_probs)
    }
/// Apply self-attention.
///
/// Return the contextualized representations and attention
/// probabilities.
Apply self-attention.
Return the contextualized representations and attention
probabilities.
Apply self - attention . Return the contextualized representations and attention probabilities .
pub fn forward_t(
        &self,
        hidden_states: &Tensor,
        attention_mask: Option<&LogitsMask>,
        train: bool,
    ) -> (Tensor, Tensor) {
        let mixed_key_layer = self.key.forward(hidden_states);
        let mixed_query_layer = self.query.forward(hidden_states);
        let mixed_value_layer = self.value.forward(hidden_states);
        let query_layer = self.transpose_for_scores(&mixed_query_layer);
        let key_layer = self.transpose_for_scores(&mixed_key_layer);
        let value_layer = self.transpose_for_scores(&mixed_value_layer);
        // Get the raw attention scores.
        let mut attention_scores = query_layer.matmul(&key_layer.transpose(-1, -2));
        attention_scores /= (self.attention_head_size as f64).sqrt();
        if let Some(mask) = attention_mask {
            attention_scores += &**mask;
        }
        // Convert the raw attention scores into a probability distribution.
        let attention_probs = attention_scores.softmax(-1, Kind::Float);
        // Drop out entire tokens to attend to, following the original
        // transformer paper.
        let attention_probs = self.dropout.forward_t(&attention_probs, train);
        let context_layer = attention_probs.matmul(&value_layer);
        let context_layer = context_layer.permute(&[0, 2, 1, 3]).contiguous();
        let mut new_context_layer_shape = context_layer.size();
        new_context_layer_shape.splice(
            new_context_layer_shape.len() - 2..,
            iter::once(self.all_head_size),
        );
        let context_layer = context_layer.view_(&new_context_layer_shape);
        (context_layer, attention_probs)
    }
pub fn forward_t ( & self , hidden_states : & Tensor , attention_mask : Option < & LogitsMask > , train : bool , ) -> ( Tensor , Tensor ) { let mixed_key_layer = self . key . forward ( hidden_states ) ; let mixed_query_layer = self . query . forward ( hidden_states ) ; let mixed_value_layer = self . value . forward ( hidden_states ) ; let query_layer = self . transpose_for_scores ( & mixed_query_layer ) ; let key_layer = self . transpose_for_scores ( & mixed_key_layer ) ; let value_layer = self . transpose_for_scores ( & mixed_value_layer ) ; // Get the raw attention scores. let mut attention_scores = query_layer . matmul ( & key_layer . transpose ( - 1 , - 2 ) ) ; attention_scores /= ( self . attention_head_size as f64 ) . sqrt ( ) ; if let Some ( mask ) = attention_mask { attention_scores += & * * mask ; } // Convert the raw attention scores into a probability distribution. let attention_probs = attention_scores . softmax ( - 1 , Kind :: Float ) ; // Drop out entire tokens to attend to, following the original // transformer paper. let attention_probs = self . dropout . forward_t ( & attention_probs , train ) ; let context_layer = attention_probs . matmul ( & value_layer ) ; let context_layer = context_layer . permute ( & [ 0 , 2 , 1 , 3 ] ) . contiguous ( ) ; let mut new_context_layer_shape = context_layer . size ( ) ; new_context_layer_shape . splice ( new_context_layer_shape . len ( ) - 2 .. , iter :: once ( self . all_head_size ) , ) ; let context_layer = context_layer . view_ ( & new_context_layer_shape ) ; ( context_layer , attention_probs ) }
Apply self-attention.
Apply self - attention .
&self
hidden_states
&Tensor
attention_mask
Option<&LogitsMask>
train
bool
&self
hidden_states
&Tensor
attention_mask
Option<&LogitsMask>
train
bool