class_docstr: Parallel chunked cross-attention layer class.
  
  Self-attention layer takes input with size [b, s, h]
  and returns output of the same size.
class_docstr_tok: ["Parallel","chunked","cross-attention","layer","class",".","Self-attention","layer","takes","input","with","size","[","b",",","s",",","h","]","and","returns","output","of","the","same","size","."]
code: def forward(
      self,
      hidden_states,
      attention_mask,
      encoder_output=None,
      set_inference_key_value_memory=False,
      inference_max_sequence_len=None,
      rotary_pos_emb=None,
      checkpoint_core_attention=False,
  ):
      if checkpoint_core_attention:
          raise ValueError(
              "checkpoint_core_attention during forward not implemented yet for ParallelChunkedCrossAttention"
          )
  
      # hidden_states is assumed to have dimension [token length, batch, dimension]
      # derive variables
      # encoder_output here is the retrieved context
      context = encoder_output
      # context is assumed to have dimension [num_chunks, num_neighbors, context_token_len, batch, dimension]
      chunk_size = self.chunk_size
      b, n, dim = (
          hidden_states.shape[1],
          hidden_states.shape[0],
          hidden_states.shape[2],
      )
      default_bias = self.cross_attention.dense.bias
      if set_inference_key_value_memory:
          seq_index = (n // chunk_size) * chunk_size
          self.current_len = n
      elif inference_max_sequence_len is not None:
          # only handles single token increment
          assert n == 1
          self.current_len += n
          chunk_id = self.current_len // chunk_size
          if chunk_id <= 0:
              # if sequence length less than chunk size, do an early return
              return torch.zeros_like(hidden_states), default_bias
          causal_padding = chunk_size - 1
          # pad it as a full chunk, put it at the end of the chunk position
          hidden_states = F.pad(
              hidden_states, (0, 0, 0, 0, causal_padding, 0), value=0.0
          )
          # only use the relevant context
          context = context[chunk_id - 1 : chunk_id, :, :, :, :]
          attention_mask = rearrange(
              attention_mask, "(b k) 1 q v -> b k 1 q v", b=b
          )
          # select the relevant chunk attn mask
          attention_mask = attention_mask[:, chunk_id - 1]
          seq_index = chunk_size
      else:
          # this is normal forward without inference
          seq_index = (n // chunk_size) * chunk_size
  
      # if sequence length less than chunk size, do an early return
      if (
          n < self.chunk_size
          and set_inference_key_value_memory
          and inference_max_sequence_len is not None
      ):
          return torch.zeros_like(hidden_states), default_bias
  
      num_chunks, num_retrieved = (
          context.shape[-5],
          context.shape[-4],
      )
  
      # causal padding
      causal_padding = chunk_size - 1
  
      x = F.pad(
          hidden_states,
          (0, 0, 0, 0, -causal_padding, causal_padding),
          value=0.0,
      )
  
      # remove sequence which is ahead of the neighbors retrieved (during inference)
  
      # seq_index = (n // chunk_size) * chunk_size
      x, x_remainder = x[:seq_index], x[seq_index:]
  
      seq_remain_len = x_remainder.shape[0]
  
      # take care of rotary positional embedding
      # make sure queries positions are properly shifted to the future
  
      if rotary_pos_emb is not None:
          q_pos_emb, k_pos_emb = rotary_pos_emb
          # currently implementation is broken
          # q need to extend to causal_padding, and just do
          # q_pos_emb = F.pad(q_pos_emb, (0, 0, -causal_padding, 0), value = 0.)
          if (
              inference_max_sequence_len is not None
              and not set_inference_key_value_memory
          ):
              token_pos = (self.current_len - 1) % chunk_size
              q_pos_emb = F.pad(
                  q_pos_emb,
                  (
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      -causal_padding - token_pos,
                      -causal_padding + token_pos,
                  ),
                  value=0.0,
              )
          else:
              q_pos_emb = F.pad(
                  q_pos_emb,
                  (0, 0, 0, 0, 0, 0, -causal_padding, 0),
                  value=0.0,
              )
  
          k_pos_emb = repeat(
              k_pos_emb, "n b h d -> (r n) b h d", r=num_retrieved
          )
          rotary_pos_emb = (q_pos_emb, k_pos_emb)
  
      # make sure number context chunks is enough
      assert x.shape[0] // chunk_size == num_chunks
  
      # reshape so we have chunk to chunk attention, without breaking causality
      x = rearrange(x, "(k n) b d -> n (b k) d", k=num_chunks)
      context = rearrange(context, "k r n b d -> (r n) (b k) d")
      # cross attention
      out, bias = self.cross_attention(
          x,
          attention_mask,
          encoder_output=context,
          rotary_pos_emb=rotary_pos_emb,
      )
  
      # reshape back to original sequence
  
      out = rearrange(out, "n (b k) d -> (k n) b d", b=b)
  
      # pad back to original, with 0s at the beginning (which will be added to the residual and be fine)
  
      out = F.pad(
          out,
          (
              0,
              0,
              0,
              0,
              causal_padding,
              -causal_padding + seq_remain_len,
          ),
          value=0.0,
      )
      if (
          not set_inference_key_value_memory
          and inference_max_sequence_len is not None
      ):
          out = out[-1:]
      return out, bias
  
imports: import math
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, LoraKQVAdapterWeightTyingConfig, LoraKVAdapterConfig, LoraQAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_softmax import MatchedScaleMaskSoftmax
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import XPOSPositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.position_embedding.rotary_position_embedding import apply_rotary_pos_emb
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, _cast_if_autocast_enabled, attention_mask_func
  from nemo.core import adapter_mixins
label_desc: Use image node_id 2 for calling the ParallelChunkedCrossAttention obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, encoder_output, set_inference_key_value_memory, inference_max_sequence_len, rotary_pos_emb, checkpoint_core_attention) and returns: out,  bias, torch,  default_bias, torch,  default_bias