Add utils functions

bfloat16_float32.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+
+import flax
+import jax
+import jax.numpy as jnp
+
+
+def to_f32(t):
+    return jax.tree_map(
+        lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t
+    )
+
+
+data = flax.serialization.msgpack_restore(
+    Path("output/flax_model.msgpack").read_bytes()
+)
+transformed_data = to_f32(data)
+Path("output/flax_model_f32.msgpack").write_bytes(
+    flax.serialization.msgpack_serialize(transformed_data)
+)

flax_to_pytorch.py

@@ -0,0 +1,4 @@
+from transformers import Wav2Vec2Model
+
+model = Wav2Vec2Model.from_pretrained("output", from_flax=True)
+model.save_pretrained("./")