params (Union[Dict, FrozenDict]) —
A PyTree of model parameters.   mask (Union[Dict, FrozenDict]) —
A PyTree with same structure as the params tree. The leaves should be booleans. It should be True
for params you want to cast, and False for those you want to skip.    Cast the floating-point params to jax.numpy.float16. This returns a new params tree and does not cast the
params in place. This method can be used on a GPU to explicitly convert th