bfloat16 support for monomer models

[oliverdutton]

Matches the casting in multimer for monomers

Runtime speedup of ~1.8x and reduces peak memory usage ~2x (for 759 residue protein, model_5)

Here's a colab link showing runtime improvement and no significant change in prediction output by visual inspection

No effect on default usage as bfloat16 turned off.

I've checked a profile on tensorboard and there's no float32 gemms with significant runtime.