def onp_fun(arg):
    if (out_dtype == jnp.bfloat16):
        return onp.trace(arg, offset, axis1, axis2, onp.float32).astype(jnp.bfloat16)
    else:
        return onp.trace(arg, offset, axis1, axis2, out_dtype)
def onp_fun(arg):
    if (out_dtype == jnp.bfloat16):
        return onp.trace(arg, offset, axis1, axis2, onp.float32).astype(jnp.bfloat16)
    else:
        return onp.trace(arg, offset, axis1, axis2, out_dtype)