def test_model_shape(
      self,
      separate_memory_values=False,
      num_intermediate_layers=None,
  ):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    config['model_config']['encoder_config'][
        'separate_memory_values'] = separate_memory_values
    config['model_config']['encoder_config'][
        'num_intermediate_layers'] = num_intermediate_layers
    config = ml_collections.FrozenConfigDict(config)

    model_config = config.model_config
    encoder_config = model_config.encoder_config

    rows = encoder_config.rows
    preprocess_fn = mention_memory_task.MentionMemoryTask.make_preprocess_fn(config)  # pylint: disable=line-too-long
    collater_fn = mention_memory_task.MentionMemoryTask.make_collater_fn(config)

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

    model = mention_memory_encoder.MentionMemoryEncoder(**encoder_config)
    dummy_input = mention_memory_task.MentionMemoryTask.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)

    memory_table = np.random.rand(rows, self.table_size // rows,
                                  encoder_config.memory_key_dim)
    memory_keys = jax.device_put_replicated(memory_table, devices)
    memory_values = memory_table.reshape(-1, encoder_config.memory_key_dim)
    memory_values = jax.device_put_replicated(memory_values, devices)
    memory_identifiers = np.arange(self.table_size)
    memory_identifiers = jax.device_put_replicated(memory_identifiers, devices)
    memory_entity_ids = memory_identifiers
    memory_text_entities = np.zeros(
        (self.table_size, encoder_config.n_memory_text_entities),
        dtype=np.int32)
    memory_text_entities = jax.device_put_replicated(memory_text_entities,
                                                     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,
        )
    initial_variables = {'params': initial_variables['params']}
    initial_variables['constants'] = {
        'memory_keys': memory_keys,
        'memory_values': memory_values,
        'memory_identifiers': memory_identifiers,
        'memory_entity_ids': memory_entity_ids,
        'memory_text_entities': memory_text_entities,
    }

    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'],
            'constants': initial_variables['constants'],
        },
        batch,
        True,
    )

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

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