def test_masking(self, retrieval_update_type, retrieval_update_config):
    """Check masked positions not contributing to input."""

    update_layer = retrieval_update_layers.RETRIEVAL_UPDATE_REGISTRY[
        retrieval_update_type](**retrieval_update_config)

    mention_mask = jnp.zeros(shape=(self.n_mentions))
    encoded_output, _ = update_layer.init_with_output(
        self.rng,
        encoded_input=self.encoded_input,
        retrieval_values=self.retrieval_values,
        retrieval_scores=self.retrieval_scores,
        mention_batch_positions=self.mention_batch_positions,
        mention_start_positions=self.mention_start_positions,
        mention_end_positions=self.mention_end_positions,
        mention_mask=mention_mask,
        deterministic=True,
    )

    if retrieval_update_type == 'dummy':
      self.assertTrue(jnp.allclose(encoded_output, self.encoded_input))
    else:
      self.assertTrue(jnp.allclose(encoded_output, self.normed_input))
Check masked positions not contributing to input.
language/mentionmemory/modules/retrieval_update_layers_test.py