@jtu.ignore_warning(category=onp.ComplexWarning)
@jtu.ignore_warning(category=RuntimeWarning, message='mean of empty slice.*')
def onp_fun(x):
    x_cast = (x if (dtype != jnp.bfloat16) else x.astype(onp.float32))
    t = (out_dtype if (out_dtype != jnp.bfloat16) else onp.float32)
    return onp_op(x_cast, axis, dtype=t, keepdims=keepdims)
@jtu.ignore_warning(category=onp.ComplexWarning)
@jtu.ignore_warning(category=RuntimeWarning, message='mean of empty slice.*')
def onp_fun(x):
    x_cast = (x if (dtype != jnp.bfloat16) else x.astype(onp.float32))
    t = (out_dtype if (out_dtype != jnp.bfloat16) else onp.float32)
    return onp_op(x_cast, axis, dtype=t, keepdims=keepdims)