def onp_fun(lhs, rhs):
    lhs = (lhs if (lhs_dtype != jnp.bfloat16) else lhs.astype(onp.float32))
    rhs = (rhs if (rhs_dtype != jnp.bfloat16) else rhs.astype(onp.float32))
    dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
    return onp.inner(lhs, rhs).astype(dtype)
def onp_fun(lhs, rhs):
    lhs = (lhs if (lhs_dtype != jnp.bfloat16) else lhs.astype(onp.float32))
    rhs = (rhs if (rhs_dtype != jnp.bfloat16) else rhs.astype(onp.float32))
    dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
    return onp.inner(lhs, rhs).astype(dtype)