def test_model_shape(
      self,
      k_top,
      num_intermediate_layers=None,
      shared_initial_encoder=True,
      shared_intermediate_encoder=True,
      shared_final_encoder=True,
      no_retrieval=False,
      same_passage_retrieval_policy='allow',
      extract_unlinked_mentions=False,
      no_retrieval_for_masked_mentions=False,
  ):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    encoder_config = copy.deepcopy(self.encoder_config)
    encoder_config['k_top'] = k_top
    encoder_config['num_intermediate_layers'] = num_intermediate_layers
    encoder_config['shared_initial_encoder'] = shared_initial_encoder
    encoder_config['shared_intermediate_encoder'] = shared_intermediate_encoder
    encoder_config['shared_final_encoder'] = shared_final_encoder
    encoder_config['no_retrieval'] = no_retrieval
    encoder_config[
        'same_passage_retrieval_policy'] = same_passage_retrieval_policy
    encoder_config['extract_unlinked_mentions'] = extract_unlinked_mentions
    encoder_config[
        'no_retrieval_for_masked_mentions'] = no_retrieval_for_masked_mentions
    config['model_config']['encoder_config'] = encoder_config
    config = ml_collections.FrozenConfigDict(config)

    model_config = config.model_config
    encoder_config = model_config.encoder_config

    preprocess_fn = readtwice_task.ReadTwiceTask.make_preprocess_fn(config)  # pylint: disable=line-too-long
    collater_fn = readtwice_task.ReadTwiceTask.make_collater_fn(config)

    test_utils.force_multi_devices(self.n_devices)
    devices = jax.local_devices()

    model = readtwice_encoder.ReadTwiceEncoder(**encoder_config)
    dummy_input = readtwice_task.ReadTwiceTask.dummy_input(config)
    dummy_input = jax.device_put_replicated(dummy_input, devices)
    init_rng = jax.random.PRNGKey(0)
    split_rng = jax.random.split(init_rng, self.n_devices)

    def model_init(*args, **kwargs):
      return model.init(*args, method=model.forward, **kwargs)

    initial_variables = jax.pmap(
        model_init, 'batch', static_broadcasted_argnums=2)(
            split_rng,
            dummy_input,
            True,
        )
    raw_example = test_utils.gen_mention_pretraining_sample(
        self.text_length,
        self.n_mentions,
        self.n_linked_mentions,
        max_length=encoder_config.max_length)
    processed_example = preprocess_fn(raw_example)
    batch = {
        key: np.tile(value, (config.per_device_batch_size, 1))
        for key, value in processed_example.items()
    }
    batch = collater_fn(batch)
    batch = {
        key: test_utils.tensor_to_numpy(value) for key, value in batch.items()
    }
    batch = {
        key: jax.device_put_replicated(value, devices)
        for key, value in batch.items()
    }

    def model_apply(*args, **kwargs):
      return model.apply(*args, method=model.forward, **kwargs)

    papply = jax.pmap(model_apply, 'batch', static_broadcasted_argnums=(2))
    encoded_output, loss_helpers, _ = papply(
        {
            'params': initial_variables['params'],
        },
        batch,
        True,
    )

    self.assertEqual(encoded_output.shape,
                     (self.n_devices, config.per_device_batch_size,
                      encoder_config.max_length, encoder_config.hidden_size))

    self.assertEqual(
        loss_helpers['target_mention_encodings'].shape,
        (self.n_devices, config.max_mention_targets *
         config.per_device_batch_size, encoder_config.memory_value_dim))
Test loss function runs and produces expected values.
language/mentionmemory/encoders/readtwice_encoder_test.py