code: def check_model_custom_n_docs(
      self,
      config,
      input_ids,
      attention_mask,
      decoder_input_ids,
      decoder_attention_mask,
      n_docs,
      **kwargs,
  ):
      self.assertIsNotNone(config.question_encoder)
      self.assertIsNotNone(config.generator)
  
      retriever = self.get_retriever(config)
  
      for model_class in self.all_model_classes:
          model = model_class(config)
          self.assertTrue(model.config.is_encoder_decoder)
  
          question_hidden_states = model.question_encoder(
              input_ids, attention_mask=attention_mask
          )[0]
  
          out = retriever(
              input_ids,
              question_hidden_states.numpy(),
              prefix=config.generator.prefix,
              return_tensors="tf",
              n_docs=n_docs,
          )
  
          (
              context_input_ids,
              context_attention_mask,
              retrieved_doc_embeds,
          ) = (
              out["context_input_ids"],
              out["context_attention_mask"],
              out["retrieved_doc_embeds"],
          )
  
          retrieved_doc_embeds = tf.cast(
              retrieved_doc_embeds, tf.float32
          )
  
          # compute doc_scores
          doc_scores = tf.squeeze(
              tf.matmul(
                  tf.expand_dims(question_hidden_states, axis=[1]),
                  retrieved_doc_embeds,
                  transpose_b=True,
              ),
              axis=[1],
          )
  
          outputs = model(
              input_ids=None,
              context_input_ids=context_input_ids,
              context_attention_mask=context_attention_mask,
              doc_scores=doc_scores,
              decoder_input_ids=decoder_input_ids,
              decoder_attention_mask=decoder_attention_mask,
              n_docs=n_docs,
          )
  
          # logits
          self.assertEqual(
              outputs.logits.shape,
              (
                  n_docs * decoder_input_ids.shape[0],
                  decoder_input_ids.shape[1],
                  config.generator.vocab_size,
              ),
          )
          # generator encoder last hidden states
          self.assertEqual(
              outputs.generator_enc_last_hidden_state.shape,
              (
                  n_docs * decoder_input_ids.shape[0],
                  self.max_combined_length,
                  config.generator.hidden_size,
              ),
          )
          # doc scores
          self.assertEqual(
              outputs.doc_scores.shape, (input_ids.shape[0], n_docs)
          )
  
imports: import json
  import os
  import shutil
  import tempfile
  import unittest
  from unittest.mock import patch
  import numpy
  from transformers import BartTokenizer
  from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES
  from transformers.models.dpr.tokenization_dpr import DPRQuestionEncoderTokenizer
  from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES
  from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
  from transformers.utils import cached_property, is_datasets_available, is_faiss_available, is_tf_available
  from ..bart.test_modeling_tf_bart import TFBartModelTester
  from ..dpr.test_modeling_tf_dpr import TFDPRModelTester
label_desc: Use image node_id 10 for calling the TFRagTestMixin obj's underlying member method code with example usage: obj.check_model_custom_n_docs(config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, n_docs) without return types