@add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
@add_code_sample_docstrings(tokenizer_class=_TOKENIZER_FOR_DOC, checkpoint=_CHECKPOINT_FOR_DOC, output_type=BaseModelOutputWithPastAndCrossAttentions, config_class=_CONFIG_FOR_DOC)
def forward(self, input_ids=None, past_key_values=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None):
    output_attentions = (output_attentions if (output_attentions is not None) else self.config.output_attentions)
    output_hidden_states = (output_hidden_states if (output_hidden_states is not None) else self.config.output_hidden_states)
    use_cache = (use_cache if (use_cache is not None) else self.config.use_cache)
    return_dict = (return_dict if (return_dict is not None) else self.config.use_return_dict)
    if ((input_ids is not None) and (inputs_embeds is not None)):
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif (input_ids is not None):
        input_shape = input_ids.size()
        input_ids = input_ids.view((- 1), input_shape[(- 1)])
        batch_size = input_ids.shape[0]
    elif (inputs_embeds is not None):
        input_shape = inputs_embeds.size()[:(- 1)]
        batch_size = inputs_embeds.shape[0]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    device = (input_ids.device if (input_ids is not None) else inputs_embeds.device)
    if (token_type_ids is not None):
        token_type_ids = token_type_ids.view((- 1), input_shape[(- 1)])
    if (position_ids is not None):
        position_ids = position_ids.view((- 1), input_shape[(- 1)])
    if (past_key_values is None):
        past_length = 0
        past_key_values = tuple(([None] * len(self.h)))
    else:
        past_length = past_key_values[0][0].size((- 2))
    device = (input_ids.device if (input_ids is not None) else inputs_embeds.device)
    if (position_ids is None):
        position_ids = torch.arange(past_length, (input_shape[(- 1)] + past_length), dtype=torch.long, device=device)
        position_ids = position_ids.unsqueeze(0).view((- 1), input_shape[(- 1)])
    if (attention_mask is not None):
        assert (batch_size > 0), 'batch_size has to be defined and > 0'
        global_attention_mask = attention_mask.view(batch_size, (- 1))
        global_attention_mask = global_attention_mask[:, None, None, :]
        global_attention_mask = global_attention_mask.to(dtype=self.dtype)
        global_attention_mask = ((1.0 - global_attention_mask) * (- 10000.0))
    else:
        global_attention_mask = None
    (batch_size, seq_length) = input_shape
    full_seq_length = (seq_length + past_length)
    head_mask = self.get_head_mask(head_mask, self.config.num_layers)
    if (inputs_embeds is None):
        inputs_embeds = self.wte(input_ids)
    position_embeds = self.wpe(position_ids)
    hidden_states = (inputs_embeds + position_embeds)
    if (token_type_ids is not None):
        token_type_embeds = self.wte(token_type_ids)
        hidden_states = (hidden_states + token_type_embeds)
    hidden_states = self.drop(hidden_states)
    output_shape = (input_shape + (hidden_states.size((- 1)),))
    presents = (() if use_cache else None)
    all_self_attentions = (() if output_attentions else None)
    all_hidden_states = (() if output_hidden_states else None)
    for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
        attn_type = self.config.attention_layers[i]
        attn_mask = global_attention_mask
        if output_hidden_states:
            all_hidden_states = (all_hidden_states + (hidden_states,))
        if (getattr(self.config, 'gradient_checkpointing', False) and self.training):
            if use_cache:
                logger.warning('`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting `use_cache=False`...')
                use_cache = False

            def create_custom_forward(module):

                def custom_forward(*inputs):
                    return module(*inputs, use_cache, output_attentions)
                return custom_forward
            outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, None, attn_mask, head_mask[i])
        else:
            try:
                outputs = block(hidden_states, layer_past=layer_past, attention_mask=attn_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions)
            except AfterStoppingPointException as e:
                raise e
            except Exception as e:
                print('failed with:')
                print(f'	 block {i}')
                print(f'	 input_ids.shape {input_ids.shape}')
                print(f'	 hidden_states.shape {hidden_states.shape}')
                print(f'	 past shapes {(layer_past[0].shape if layer_past else layer_past)}')
                raise e
        hidden_states = outputs[0]
        if (use_cache is True):
            presents = (presents + (outputs[1],))
        if output_attentions:
            all_self_attentions = (all_self_attentions + (outputs[(2 if use_cache else 1)],))
    hidden_states = self.ln_f(hidden_states)
    hidden_states = hidden_states.view(*output_shape)
    if output_hidden_states:
        all_hidden_states = (all_hidden_states + (hidden_states,))
    if (not return_dict):
        return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if (v is not None)))
    return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
@add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
@add_code_sample_docstrings(tokenizer_class=_TOKENIZER_FOR_DOC, checkpoint=_CHECKPOINT_FOR_DOC, output_type=BaseModelOutputWithPastAndCrossAttentions, config_class=_CONFIG_FOR_DOC)
def forward(self, input_ids=None, past_key_values=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None):
    output_attentions = (output_attentions if (output_attentions is not None) else self.config.output_attentions)
    output_hidden_states = (output_hidden_states if (output_hidden_states is not None) else self.config.output_hidden_states)
    use_cache = (use_cache if (use_cache is not None) else self.config.use_cache)
    return_dict = (return_dict if (return_dict is not None) else self.config.use_return_dict)
    if ((input_ids is not None) and (inputs_embeds is not None)):
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif (input_ids is not None):
        input_shape = input_ids.size()
        input_ids = input_ids.view((- 1), input_shape[(- 1)])
        batch_size = input_ids.shape[0]
    elif (inputs_embeds is not None):
        input_shape = inputs_embeds.size()[:(- 1)]
        batch_size = inputs_embeds.shape[0]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    device = (input_ids.device if (input_ids is not None) else inputs_embeds.device)
    if (token_type_ids is not None):
        token_type_ids = token_type_ids.view((- 1), input_shape[(- 1)])
    if (position_ids is not None):
        position_ids = position_ids.view((- 1), input_shape[(- 1)])
    if (past_key_values is None):
        past_length = 0
        past_key_values = tuple(([None] * len(self.h)))
    else:
        past_length = past_key_values[0][0].size((- 2))
    device = (input_ids.device if (input_ids is not None) else inputs_embeds.device)
    if (position_ids is None):
        position_ids = torch.arange(past_length, (input_shape[(- 1)] + past_length), dtype=torch.long, device=device)
        position_ids = position_ids.unsqueeze(0).view((- 1), input_shape[(- 1)])
    if (attention_mask is not None):
        assert (batch_size > 0), 'batch_size has to be defined and > 0'
        global_attention_mask = attention_mask.view(batch_size, (- 1))
        global_attention_mask = global_attention_mask[:, None, None, :]
        global_attention_mask = global_attention_mask.to(dtype=self.dtype)
        global_attention_mask = ((1.0 - global_attention_mask) * (- 10000.0))
    else:
        global_attention_mask = None
    (batch_size, seq_length) = input_shape
    full_seq_length = (seq_length + past_length)
    head_mask = self.get_head_mask(head_mask, self.config.num_layers)
    if (inputs_embeds is None):
        inputs_embeds = self.wte(input_ids)
    position_embeds = self.wpe(position_ids)
    hidden_states = (inputs_embeds + position_embeds)
    if (token_type_ids is not None):
        token_type_embeds = self.wte(token_type_ids)
        hidden_states = (hidden_states + token_type_embeds)
    hidden_states = self.drop(hidden_states)
    output_shape = (input_shape + (hidden_states.size((- 1)),))
    presents = (() if use_cache else None)
    all_self_attentions = (() if output_attentions else None)
    all_hidden_states = (() if output_hidden_states else None)
    for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
        attn_type = self.config.attention_layers[i]
        attn_mask = global_attention_mask
        if output_hidden_states:
            all_hidden_states = (all_hidden_states + (hidden_states,))
        if (getattr(self.config, 'gradient_checkpointing', False) and self.training):
            if use_cache:
                logger.warning('`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting `use_cache=False`...')
                use_cache = False

            def create_custom_forward(module):

                def custom_forward(*inputs):
                    return module(*inputs, use_cache, output_attentions)
                return custom_forward
            outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, None, attn_mask, head_mask[i])
        else:
            try:
                outputs = block(hidden_states, layer_past=layer_past, attention_mask=attn_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions)
            except AfterStoppingPointException as e:
                raise e
            except Exception as e:
                print('failed with:')
                print(f'	 block {i}')
                print(f'	 input_ids.shape {input_ids.shape}')
                print(f'	 hidden_states.shape {hidden_states.shape}')
                print(f'	 past shapes {(layer_past[0].shape if layer_past else layer_past)}')
                raise e
        hidden_states = outputs[0]
        if (use_cache is True):
            presents = (presents + (outputs[1],))
        if output_attentions:
            all_self_attentions = (all_self_attentions + (outputs[(2 if use_cache else 1)],))
    hidden_states = self.ln_f(hidden_states)
    hidden_states = hidden_states.view(*output_shape)
    if output_hidden_states:
        all_hidden_states = (all_hidden_states + (hidden_states,))
    if (not return_dict):
        return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if (v is not None)))
    return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)