code: def forward(
      self,
      input: torch.Tensor,  # (batch, seqlen) int32
      input_sub: torch.Tensor,  # (batch, seqlen) int32
      length: torch.Tensor,  # (batch) int32
      context: torch.Tensor,  # (batch, seqlen) bool
      sample_ids: torch.Tensor,  # (batch, seq_len) int32
      num_segments: torch.Tensor,  # (batch, seq_len) int32
      segment: torch.Tensor,  # (batch, seqlen) int32
      segment_rel_offset: torch.Tensor,  # (batch, seq_len) int32
      segment_rel: torch.Tensor,  # (batch, num_segment_bucket) int32
      span: torch.Tensor,  # (batch, seqlen) int32
      ext_table_ids: torch.Tensor,  # (ext_table_size) int32
      ext_table_sub: torch.Tensor,  # (ext_table_size) int32
  ):
      batch = input.size(0)
      seqlen = input.size(1)
      # processing masks and position bias bucket
      with torch.no_grad():
          device = input.device
  
          # calc segment bucket
          segment_rel_2d = torch.masked_fill(
              segment[:, :, None] * num_segments[:, :, None]
              + segment[:, None, :]
              + segment_rel_offset[:, :, None],
              ~(
                  (sample_ids[:, :, None] == sample_ids[:, None, :])
                  & (span[:, None, :] == span[:, :, None])
              ),  # not in the same span or sample
              0,  # avoid torch.gather overflow
          ).view(batch, seqlen * seqlen)
  
          segment_bucket = torch.gather(
              input=segment_rel,
              dim=1,
              index=segment_rel_2d.long(),
          ).view(batch, seqlen, seqlen)
  
          segment_bucket.masked_fill_(
              ~(
                  (sample_ids[:, :, None] == sample_ids[:, None, :])
                  & (span[:, None, :] == span[:, :, None])
              ),  # not in the same span or sample
              1,  # bucket is used for in-context samples
          )
  
          # directional mask
          directional_mask_2d = torch.arange(
              seqlen, device=device
          ) <= torch.arange(seqlen, device=device).view(-1, 1)
          # sample mask
          sample_mask_2d = (sample_ids[:, :, None] == 0) | (
              sample_ids[:, :, None] == sample_ids[:, None, :]
          )
          # context mask
          attention_mask = context[:, None, :] | (
              context[:, :, None].logical_not()
              & directional_mask_2d.view(1, seqlen, seqlen)
          )
          # span mask
          attention_mask = (
              attention_mask
              & sample_mask_2d
              & (span[:, None, :] == span[:, :, None])
          )
          # length mask
          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 = torch.arange(seqlen, device=device).expand(
              batch, seqlen
          )
  
      hidden_states = self.input_embedding(input, input_sub)
      position_bias = self.position_bias(
          position, position, segment_bucket
      )
  
      hidden_states = self.encoder(
          hidden_states, attention_mask, position_bias
      )
  
      ext_table = self.input_embedding(ext_table_ids, ext_table_sub)
  
      logits = self.input_embedding.projection(hidden_states, ext_table)
      return logits, hidden_states
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/models/bee_torch.py
imports: from typing import Optional, Tuple
  import torch
  from ..native_layers import Encoder, EmbeddingExt, BucketPositionBias
  from .bee import CPMBeeConfig, CPMBeeInferenceState
label_desc: Use image node_id 2 for calling the CPMBeeTorch obj's underlying member method code with example usage: obj.forward(input, input_sub, length, context, sample_ids, num_segments, segment, segment_rel_offset, segment_rel, span, ext_table_ids, ext_table_sub) and returns: logits,  hidden_states