结合使用Cupy-Numpy的优雅方式

Question

我正在构建一个Python包，其中的定义/方法正在使用Numpy。我想让我的未来用户在加载程序包时选择Cupy或Numpy。高雅的说法还是让我说出正确的做法？

提供更多背景信息；我的图书馆位于一个名为mynewlibrary的文件夹中，并且该文件夹中只有一个文件； __init__.py。目前，在__init__.py内，我都按如下所示进行导入。

import numpy as np_cpu
import cupy as np

def compute(a):
    return np.pi*a

如何在import mynewlibrary期间按需将np设置为numpy或cupy？

Answer 1

一种幼稚的方法是使用包含布尔变量的虚拟模块config.py：

# config.py
use_numpy=True

在您的__init__.py文件中，导入config.py并检查use_numpy的值：

import config

if config.use_numpy:
   import numpy as np
else:
   import cupy as np

def compute(a):
    return np.pi*a

如果要使用cupy，请在use_numpy=False中设置config.py。