Restrict bfloat16 conversion to only float32 types.

PiperOrigin-RevId: 868479414

Author: jerryxyj

export/orbax/export/utils.py

@@ -544,7 +544,7 @@ def to_bfloat16(x: Any) -> Any:
 
   def _to_bfloat16_leaf(x: Any) -> Any:
     if isinstance(x, jax.ShapeDtypeStruct):
-      if jnp.issubdtype(x.dtype, jnp.floating):
+      if jnp.issubdtype(x.dtype, jnp.float32):
         return jax.ShapeDtypeStruct(
             x.shape,
             jnp.bfloat16,
@@ -555,12 +555,8 @@ def _to_bfloat16_leaf(x: Any) -> Any:
     if hasattr(x, 'dtype'):
       if x.dtype == tf.string:
         return x
-      if jnp.issubdtype(x.dtype, jnp.floating):
+      if jnp.issubdtype(x.dtype, jnp.float32):
         return x.astype(jnp.bfloat16)
-
-    if isinstance(x, float):
-      return jnp.bfloat16(x)
-
     return x
 
   flattened_x, treedef = jax.tree_util.tree_flatten(x)

export/orbax/export/utils_test.py

@@ -306,7 +306,7 @@ def test_to_bfloat16(self):
     self.assertEqual(y['a'].dtype, jnp.bfloat16)
     self.assertEqual(y['b'].dtype, tf.string)
     self.assertAllEqual(y['b'], x['b'])
-    self.assertEqual(y['c'].dtype, jnp.bfloat16)
+    self.assertIsInstance(y['c'], float)
     self.assertEqual(y['d'].dtype, jnp.bfloat16)
 
 if __name__ == '__main__':