def onp_fun(x, y):
    dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
    return onp.matmul(x, y).astype(dtype)
def onp_fun(x, y):
    dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
    return onp.matmul(x, y).astype(dtype)