src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py/GPTNeoXJapaneseLayer/forward
class GPTNeoXJapaneseLayer:    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        use_cache=False,
        layer_past=None,
        output_attentions=False,
    ):
        residual = hidden_states
        ln_out = self.input_layernorm(hidden_states)
        attention_layer_outputs, attn_bias = self.attention(
            ln_out,
            attention_mask=attention_mask,
            layer_past=layer_past,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        attn_output = attention_layer_outputs[0]  # output_attn: a, present, (attentions)
        outputs = attention_layer_outputs[1:]

        # attn_output = (atten_output + bias) + residual
        attn_output = bias_dropout_add(
            attn_output,
            bias=attn_bias.expand_as(residual) if attn_bias is not None else attn_bias,
            residual=residual,
            prob=self.hidden_dropout,
            training=self.training,
        )
        mlp_output = self.mlp(self.post_attention_layernorm(attn_output))

        # attn_output = (mlp_output + mlp_bias) + atten_output
        attn_output = bias_dropout_add(
            mlp_output, bias=None, residual=attn_output, prob=self.hidden_dropout, training=self.training
        )

        if use_cache:
            outputs = (attn_output,) + outputs
        else:
            outputs = (attn_output,) + outputs[1:]

        return outputs  # hidden_states, present, (attentions)