def test_saving_loading_weights_returns_original(self):
    """Test whether saving weights and then loading them returns original."""

    # Create initial variables
    init_rng = jax.random.PRNGKey(self.config.seed)

    dummy_input = example_task.ExampleTask.dummy_input(self.config)

    model = example_task.ExampleTask.build_model(self.model_config)
    initial_variables = jax.jit(model.init)(init_rng, dummy_input, True)

    # Save weights
    checkpoint_utils.save_weights(self.weight_path, initial_variables['params'])

    # Load weights
    loaded_params = checkpoint_utils.load_weights(self.weight_path)

    arrayeq = lambda x, y: jnp.all(x == y)
    self.assertTrue(
        jax.tree.map(
            arrayeq,
            loaded_params,
            flax.core.unfreeze(initial_variables['params']),
        )
    )
Test whether saving weights and then loading them returns original.
language/mentionmemory/utils/checkpoint_utils_test.py