code: def forward(
      self,
      input: torch.Tensor,  # (batch, seqlen)
      length: torch.Tensor,  # (batch)
      context: torch.Tensor,  # (batch, seqlen)
      position: torch.Tensor,  # (batch, seqlen)
      segment: torch.Tensor,  # (batch, seqlen)
      span: torch.Tensor,  # (batch, seqlen)
  ):
      batch = input.size(0)
      seqlen = input.size(1)
      input_prompt = input[:, : self.prompt_length].contiguous()
      input_ids = input[:, self.prompt_length :].contiguous()
  
      prompt_states = self.prompt_embedding(input_prompt)
      hidden_states = self.input_embedding(input_ids)
      segment_states = self.segment_embedding(segment)
      hidden_states = (
          torch.cat([prompt_states, hidden_states], 1) + segment_states
      )
  
      with torch.no_grad():
          device = input.device
          directional_mask_2d = torch.arange(
              seqlen, device=device
          ) <= torch.arange(seqlen, device=device).view(-1, 1)
          attention_mask = context[:, None, :] | (
              context[:, :, None].logical_not()
              & directional_mask_2d.view(1, seqlen, seqlen)
          )
          attention_mask = attention_mask & (
              span[:, None, :] == span[:, :, None]
          )
          mask_1d = (
              torch.arange(seqlen, device=device)[None, :].repeat(
                  batch, 1
              )
              < length[:, None]
          )
          attention_mask = (
              mask_1d.view(batch, seqlen, 1)
              & mask_1d.view(batch, 1, seqlen)
              & attention_mask
          )
  
      position_bias = self.position_bias(
          position, position, segment, segment
      )
      hidden_states = self.encoder(
          hidden_states, attention_mask, position_bias
      )
  
      logits = self.input_embedding.projection(hidden_states)
      return logits, hidden_states
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/models/ant_torch.py
imports: import torch
  from ..native_layers import Encoder, Embedding, SegmentPositionEmbedding
  from .ant import CPMAntConfig
label_desc: Use image node_id 2 for calling the CPMAntTorch obj's underlying member method code with example usage: obj.forward(input, length, context, position, segment, span) and returns: logits,  hidden_states