   Cast the floating-point params to jax.numpy.bfloat16. This returns a new params tree and does not cast
the params in place. This method can be used on a TPU to explicitly convert the model parameters to bfloat16 precision to do full
half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed.  Examples:   Copied >>> from diffusers import FlaxUNet2DConditionModel

>>> # load model
>>> model, 