No public description

PiperOrigin-RevId: 865066915

Author: PAenugula

model/orbax/experimental/model/jd2obm/utils.py

@@ -25,13 +25,17 @@
 
 def _obm_to_voxel_dtype(t):
   if isinstance(t, obm.ShloDType):
+    if t == obm.ShloDType.str:
+      return np.dtype(np.str_)
     return obm.shlo_dtype_to_np_dtype(t)
   return t
 
 
 def _voxel_to_obm_dtype(t) -> obm.ShloDType:
   if not isinstance(t, np.dtype):
     raise ValueError(f'Expected a numpy.dtype, got {t!r} of type {type(t)}')
+  if t == np.dtype(np.str_) or t == np.dtype(np.bytes_):
+    return obm.ShloDType.str
   return obm.np_dtype_to_shlo_dtype(t)