PyTree with same structure as the params tree. The leaves should be booleans, True for params
you want to cast, and should be False for those you want to skip.



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 TPU to explicitly convert the model parameters to bfloat16 precision to do full
half-precision training or to save weights in bfloat16