def run_model(self, config, entity_vocab_size):
    """Initialize and run the model once, perform sanity checks."""
    np.random.seed(0)

    # Save arrays to test retrieval saver.
    memory_identifiers = np.arange(self.table_size)
    memory_identifiers = jax.device_put_replicated(memory_identifiers,
                                                   self.devices)
    memory_entity_ids = memory_identifiers
    config['memory_entity_id_pattern'] = self.save_sharded_array(
        memory_entity_ids, 'memory_entity_id')
    memory_text = np.random.randint(
        config['model_config']['encoder_config']['vocab_size'],
        size=(self.n_devices, self.table_size, self.memory_text_length),
        dtype=np.int32)
    config['memory_text_pattern'] = self.save_sharded_array(
        memory_text, 'memory_text')
    memory_positions = np.random.randint(
        self.memory_text_length,
        size=(self.n_devices, self.table_size, 2),
        dtype=np.int32)
    config['memory_positions_pattern'] = self.save_sharded_array(
        memory_positions, 'memory_positions')

    config = ml_collections.FrozenConfigDict(config)
    model_config = config.model_config
    encoder_config = model_config.encoder_config

    rows = encoder_config.rows
    preprocess_fn = mention_based_entity_qa_task.MentionBasedEntityQATask.make_preprocess_fn(config)  # pylint: disable=line-too-long
    collater_fn = mention_based_entity_qa_task.MentionBasedEntityQATask.make_collater_fn(
        config)
    postprocess_fn = mention_based_entity_qa_task.MentionBasedEntityQATask.make_output_postprocess_fn(
        config)

    model = mention_based_entity_qa_task.MentionBasedEntityQATask.build_model(
        model_config)
    dummy_input = mention_based_entity_qa_task.MentionBasedEntityQATask.dummy_input(
        config)
    dummy_input = jax.device_put_replicated(dummy_input, self.devices)
    init_rng = jax.random.PRNGKey(0)
    split_rng = jax.random.split(init_rng, self.n_devices)

    memory_table = np.random.rand(rows, self.table_size // rows,
                                  encoder_config.memory_key_dim)
    memory_keys = jax.device_put_replicated(memory_table, self.devices)
    memory_values = memory_table.reshape(-1, encoder_config.memory_key_dim)
    memory_values = jax.device_put_replicated(memory_values, self.devices)

    initial_variables = jax.pmap(
        model.init, 'batch', static_broadcasted_argnums=2)(
            split_rng,
            dummy_input,
            True,
        )
    initial_variables = {'params': initial_variables['params']}
    initial_variables['constants'] = {
        'encoder': {
            'memory_keys': memory_keys,
            'memory_values': memory_values,
            'memory_identifiers': memory_identifiers,
            'memory_entity_ids': memory_entity_ids,
        }
    }

    def sample_batch():
      processed_examples = []
      for _ in range(config.per_device_batch_size):
        raw_example = test_utils.gen_mention_pretraining_sample(
            self.text_length,
            self.n_mentions,
            self.n_linked_mentions,
            entity_vocab_size=entity_vocab_size,
            max_length=encoder_config.max_length)
        processed_example = preprocess_fn(raw_example)
        processed_examples.append(processed_example)
      batch = stack(processed_examples)
      batch = collater_fn(batch)
      batch = {
          key: test_utils.tensor_to_numpy(value)
          for key, value in batch.items()
      }
      return batch

    batch = stack([sample_batch() for _ in range(self.n_devices)])
    batch = {
        key: jax.device_put_sharded(list(value), self.devices)
        for key, value in batch.items()
    }

    loss_fn = jax.pmap(
        mention_based_entity_qa_task.MentionBasedEntityQATask.make_loss_fn(
            config),
        'batch',
        static_broadcasted_argnums=(0, 4))
    _, metrics, auxiliary_output = loss_fn(
        model_config,
        initial_variables['params'],
        {'constants': initial_variables['constants']},
        batch,
        True,
    )

    self.assertArrayEqual(metrics['agg']['denominator'],
                          batch['mention_target_weights'].sum(1))

    features = postprocess_fn(batch, auxiliary_output)
    # Check features are JSON-serializable
    json.dumps(features)
    # Check features match the original batch
    for key in batch.keys():
      self.assertArrayEqual(np.array(features[key]), batch[key])

    n_mentions_per_device = (
        config.per_device_batch_size * config.max_mention_targets)
    k_top_final = (
        encoder_config.final_k_top_post_selection or
        encoder_config.final_k_top_device * self.n_devices)
    self.assertSequenceEqual(
        np.array(features['memory_text']).shape, [
            self.n_devices, n_mentions_per_device, k_top_final,
            self.memory_text_length
        ])
    self.assertSequenceEqual(
        np.array(features['memory_positions']).shape,
        [self.n_devices, n_mentions_per_device, k_top_final, 2])

    return batch, initial_variables, metrics
Initialize and run the model once, perform sanity checks.
language/mentionmemory/tasks/mention_based_entity_qa_task_test.py