class_docstr: This class serves as a micro library for forward function substitution of GPTJ models
  under pipeline setting.
docstr: labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
      Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
      config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
      `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
  
  # This function is modified on the basis of transformers.models.gptj.modeling_gptj.GPTJForSequenceClassification.forward.
  # Please refer to original code of transformers for more details.
docstr_tok: ["labels","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",",")","`",",","*","optional","*",")",":","Labels","for","computing","the","sequence","classification\/regression","loss",".","Indices","should","be","in","`","[","0",",","...",",","config.num_labels","-","1","]","`",".","If","`","config.num_labels","==","1","`","a","regression","loss","is","computed","(","Mean-Square","loss",")",",","If","`","config.num_labels",">","1","`","a","classification","loss","is","computed","(","Cross-Entropy",")",".","#","This","function","is","modified","on","the","basis","of","transformers.models.gptj.modeling_gptj.GPTJForSequenceClassification.forward",".","#","Please","refer","to","original","code","of","transformers","for","more","details","."]
code: def gptj_for_sequence_classification_forward(
      self: GPTJForSequenceClassification,
      input_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
      attention_mask: Optional[torch.FloatTensor] = None,
      token_type_ids: Optional[torch.LongTensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      head_mask: Optional[torch.FloatTensor] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      labels: Optional[torch.LongTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      stage_manager: Optional[PipelineStageManager] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_index: Optional[List[int]] = None,
      shard_config: ShardConfig = None,
  ) -> Union[Dict, Tuple, SequenceClassifierOutputWithPast]:
      r"""
      labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
          Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
          config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
          `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
  
      # This function is modified on the basis of transformers.models.gptj.modeling_gptj.GPTJForSequenceClassification.forward.
      # Please refer to original code of transformers for more details.
      """
      logger = logging.get_logger(__name__)
  
      if input_ids is not None:
          batch_size, _ = input_ids.shape[:2]
      else:
          batch_size, _ = hidden_states.shape[:2]
      assert (
          self.config.pad_token_id is not None or batch_size == 1
      ), "Cannot handle batch sizes > 1 if no padding token is defined."
  
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      transformer_outputs = GPTJPipelineForwards.gptj_model_forward(
          self.transformer,
          input_ids,
          past_key_values=past_key_values,
          attention_mask=attention_mask,
          token_type_ids=token_type_ids,
          position_ids=position_ids,
          head_mask=head_mask,
          inputs_embeds=inputs_embeds,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          stage_manager=stage_manager,
          hidden_states=hidden_states,
          stage_index=stage_index,
          shard_config=shard_config,
      )
  
      # If not at the last stage, return hidden_states as in GPTJModel
      if not stage_manager.is_last_stage():
          return {"hidden_states": transformer_outputs["hidden_states"]}
  
      hidden_states = transformer_outputs[0]
      logits = self.score(hidden_states)
  
      if self.config.pad_token_id is None:
          sequence_lengths = -1
      else:
          if input_ids is not None:
              sequence_lengths = (
                  torch.ne(input_ids, self.config.pad_token_id).sum(-1)
                  - 1
              ).to(logits.device)
          else:
              sequence_lengths = -1
              logger.warning_once(
                  f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
                  "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
              )
  
      pooled_logits = logits[
          torch.arange(batch_size, device=logits.device),
          sequence_lengths,
      ]
  
      loss = None
      if labels is not None:
          labels = labels.to(pooled_logits.device)
          if self.config.problem_type is None:
              if self.num_labels == 1:
                  self.config.problem_type = "regression"
              elif self.num_labels > 1 and (
                  labels.dtype == torch.long
                  or labels.dtype == torch.int
              ):
                  self.config.problem_type = (
                      "single_label_classification"
                  )
              else:
                  self.config.problem_type = (
                      "multi_label_classification"
                  )
  
          if self.config.problem_type == "regression":
              loss_fct = MSELoss()
              if self.num_labels == 1:
                  loss = loss_fct(
                      pooled_logits.squeeze(), labels.squeeze()
                  )
              else:
                  loss = loss_fct(pooled_logits, labels)
          elif (
              self.config.problem_type == "single_label_classification"
          ):
              loss_fct = CrossEntropyLoss()
              loss = loss_fct(
                  pooled_logits.view(-1, self.num_labels),
                  labels.view(-1),
              )
          elif self.config.problem_type == "multi_label_classification":
              loss_fct = BCEWithLogitsLoss()
              loss = loss_fct(pooled_logits, labels)
      if not return_dict:
          output = (pooled_logits,) + transformer_outputs[1:]
          return ((loss,) + output) if loss is not None else output
  
      return SequenceClassifierOutputWithPast(
          loss=loss,
          logits=pooled_logits,
          past_key_values=transformer_outputs.past_key_values,
          hidden_states=transformer_outputs.hidden_states,
          attentions=transformer_outputs.attentions,
      )
  
file_path: colossalai/colossalai/shardformer/modeling/gptj.py
imports: from typing import Dict, List, Optional, Tuple, Union
  import torch
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast
  from transformers.models.gptj.modeling_gptj import GPTJForCausalLM, GPTJForQuestionAnswering, GPTJForSequenceClassification, GPTJModel, apply_rotary_pos_emb, get_embed_positions
  from transformers.utils import is_torch_fx_proxy, logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
  from colossalai.shardformer.shard import ShardConfig
label_desc: Use image node_id 3 for calling the GPTJPipelineForwards obj's underlying member method code with example usage: obj.gptj_for_sequence_classification_forward(input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, stage_manager, hidden_states, stage_index, shard_config) and returns: SequenceClassifierOutputWithPast, dict, unknown