Load pretrained jax weights

[l-bick]

Added option to load pretrained jax-weights (such as from Alphafold) to be loaded for training. Implementation is analog to run_pretrained_openfold.

[gahdritz]

This is a nice option. One quibble: it looks like the "help" field is wrong ("Path to JAX model parameters. If None, and openfold_checkpoint_path is also None, parameters are selected automatically according to the model name from openfold/resources/params"). It doesn't look like pretrained parameters are automatically loaded if no checkpoint is specified, and it should stay that way; by default, training should obviously commence from a random initialization. I think we should change the option to "--resume_from_jax_params" to match the corresponding OF flag and change the help message to "Path to an .npz JAX parameter file with which to initialize the model"

[l-bick]

You are right, thanks for pointing it out. The help message was uncarefully taken from run_pretrained. Ive adapted the conventions as suggested.

[gahdritz]

Thanks for the PR!