def assign_weights(model, checkpoint_reader,
                       variable_to_checkpoint_name_mapping):
      """Assign weights from checkpoint to model."""
      for variable_name, ckpt_variable_name in (
          variable_to_checkpoint_name_mapping.items()):
        variable = get_variable_by_name(model, variable_name)
        ckpt_weight = checkpoint_reader.get_tensor(ckpt_variable_name)
        if ckpt_weight.shape != variable.shape:
          logging.warning('Shape doesnt match for %s: %s vs %s', variable_name,
                          variable.shape, ckpt_weight.shape)
          continue
        variable.assign(ckpt_weight)
Assign weights from checkpoint to model.
language/search_agents/muzero/transformer_encoder.py