Allow pipeline to run in bfloat16

[pcuenca]

Changes:

• Allow dtype to be specified on model load. This is a temporary solution until Save training dtype as part of the configuration #567 is addressed in a more principled way.
• Convert loaded params to bfloat16 or float16 if necessary.
• Keep latents in float32 during inference loop.

See comment: #559 (comment)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[patil-suraj]

The dtype here specifies the dtype of computation and should not be used for params. We should actually add a different argument to specify the dtype of parameters. For example, flax now has a param_dtype argument https://github.com/google/flax/blob/main/flax/linen/linear.py#L78.

The reason to do this is that, it's not recommended to keep params in bfloat16 on TPUs during training as it can cause numerical instabilities, it's only done for extremely large models. Also having these two arguments allows you to do mixed-precision training: keep params in fp32 but do forward computation in half precision.

For now we could create a branch with bfloat16 weights and the use that for inference.

[pcuenca]

I removed the automatic conversion but left the other change in order to be able to load the models specifying a dtype. @patil-suraj please, let me know if that's ok.

[patil-suraj]

LGTM, thanks a lot!

[patil-suraj]

This file doesn't seem to be in main. Should be removed from here before mergin

[pcuenca]

The merge target is flax_pipeline for now.

[pcuenca]

Replaced by #600.