src/transformers/models/owlvit/convert_owlvit_original_flax_to_hf.py/to_f32
def to_f32(params):
    return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, params)