def to_tf_dtype(jax_dtype):
    if (jax_dtype == dtypes.float0):
        jax_dtype = dtypes.bfloat16
    return tf.dtypes.as_dtype(jax_dtype)
def to_tf_dtype(jax_dtype):
    if (jax_dtype == dtypes.float0):
        jax_dtype = dtypes.bfloat16
    return tf.dtypes.as_dtype(jax_dtype)