def onp_fun(*args):
    args = [(x if (x.dtype != jnp.bfloat16) else x.astype(onp.float32)) for x in args]
    dtype = functools.reduce(jnp.promote_types, arg_dtypes)
    return onp.concatenate(args, axis=axis).astype(dtype)
def onp_fun(*args):
    args = [(x if (x.dtype != jnp.bfloat16) else x.astype(onp.float32)) for x in args]
    dtype = functools.reduce(jnp.promote_types, arg_dtypes)
    return onp.concatenate(args, axis=axis).astype(dtype)