Figure out how to deal with `keep_cupy_as_array` argument

[flying-sheep]

Should we make it public or change the API to return an array by default?