def test_load_weights(self, separate_memory_values=False, memory_only=False):
    """Test saving and loading model recovers original parameters."""

    config = copy.deepcopy(self.config)
    config['model_config']['encoder_config'][
        'separate_memory_values'] = separate_memory_values
    config = ml_collections.ConfigDict(config)

    model_config = config.model_config
    encoder_config = model_config.encoder_config
    rows = encoder_config.rows
    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,
    }
    n_shards = 4

    tempdir_obj = self.create_tempdir()
    tempdir = tempdir_obj.full_path

    memory_key_base = os.path.join(tempdir, 'memory_keys')
    memory_value_base = os.path.join(tempdir, 'memory_values')
    memory_id_base = os.path.join(tempdir, 'memory_id')
    memory_entity_id_base = os.path.join(tempdir, 'memory_entity_id')
    memory_text_entities_base = os.path.join(tempdir, 'memory_text_entities')

    unreplicated_variables = jax_utils.unreplicate(initial_variables)
    unreplicated_variables['params'] = flax.core.unfreeze(
        unreplicated_variables['params']
    )

    if memory_only:
      load_weights = 'memory_only'
    else:
      load_weights = os.path.join(tempdir, 'weights')
      checkpoint_utils.save_weights(
          load_weights, unreplicated_variables['params']
      )

    memory_keys = initial_variables['constants']['memory_keys']
    memory_keys = memory_keys.reshape(
        n_shards, -1, encoder_config.memory_key_dim
    )
    memory_values = initial_variables['constants']['memory_values']
    memory_values = memory_values.reshape(
        n_shards, -1, encoder_config.memory_key_dim
    )
    memory_ids = initial_variables['constants']['memory_identifiers'].reshape(
        n_shards, -1
    )
    memory_entity_ids = initial_variables['constants'][
        'memory_entity_ids'
    ].reshape(n_shards, -1)
    memory_text_entities = initial_variables['constants'][
        'memory_text_entities'
    ].reshape(n_shards, -1, encoder_config.n_memory_text_entities)

    for shard in range(n_shards):
      np.save(memory_key_base + str(shard), memory_keys[shard])
      np.save(memory_value_base + str(shard), memory_values[shard])
      np.save(memory_id_base + str(shard), memory_ids[shard])
      np.save(memory_entity_id_base + str(shard), memory_entity_ids[shard])
      np.save(memory_entity_id_base + str(shard), memory_entity_ids[shard])
      np.save(memory_text_entities_base + str(shard),
              memory_text_entities[shard])

    config.memory_key_pattern = memory_key_base + '*'
    config.memory_value_pattern = memory_value_base + '*'
    config.memory_id_pattern = memory_id_base + '*'
    config.memory_entity_id_pattern = memory_entity_id_base + '*'
    config.memory_text_entities_pattern = memory_text_entities_base + '*'
    config.load_weights = load_weights

    loaded_variables = mention_memory_encoder.MentionMemoryEncoder.load_weights(
        config)

    arrayeq = lambda x, y: jnp.all(x == y)
    constants = {
        key: value
        for key, value in initial_variables['constants'].items()
        if not (key == 'memory_values' and not separate_memory_values)
    }
    comparison_variables = {'constants': constants}
    if not memory_only:
      comparison_variables['params'] = flax.core.unfreeze(
          initial_variables['params']
      )

    self.assertTrue(
        jax.tree.map(arrayeq, loaded_variables, comparison_variables)
    )
Test saving and loading model recovers original parameters.
language/mentionmemory/encoders/mention_memory_encoder_test.py