pybind11とxtensorによるPythonとC++の連携

前回の記事でNumCppを試しましたが、全ての配列を2次元配列で扱う仕様がいまいちだったのでEigenを使ったほうが良さそうだという結論になりました。 kivantium.hateblo.jp

この記事ではxtensorという別のライブラリを試してみます。xtensorの開発はQuantStackという会社が中心になって行われているようです。Webサイトを見る限りJupyterLabなどのOSSを開発している会社のようですが、どうやって儲けているのかいまいち分かりませんでした……

インストール

READMEを見るとcondaでのインストールが推奨されているのですが、header-onlyライブラリなのでソースコードを落とすだけで使うことが出来ます。Pythonバインディングを使うためには3つのプロジェクトが必要になります。ここでは作業ディレクトリで以下のコマンドを実行した場合を説明します。

git clone https://github.com/xtensor-stack/xtensor.git
git clone https://github.com/xtensor-stack/xtl.git
git clone https://github.com/xtensor-stack/xtensor-python.git

簡単な例

前回と同様に足し算の例で試してみます。以下のC++コードを example.cpp という名前で保存します。

#include <pybind11/pybind11.h>
#include <xtensor/xmath.hpp>
#define FORCE_IMPORT_ARRAY
#include <xtensor-python/pyarray.hpp>

xt::pyarray<double> add_arrays(xt::pyarray<double>& input1, xt::pyarray<double>& input2) {
    return input1 + input2;
}

PYBIND11_MODULE(example, m)
{
    xt::import_numpy();
    m.def("add_arrays", &add_arrays);
}

コンパイルコマンドは以下の通りです。長いですが、pybind11のコンパイルコマンドをC++14向けに変更してxtensorのディレクトリを追加しただけです。

c++ -O3 -Wall -shared -std=c++14 -fPIC $(python3 -m pybind11 --includes) -I xtensor/include/ -I xtl/include/ -I xtensor-python/include/ -DNUMCPP_INCLUDE_PYBIND_PYTHON_INTERFACE example.cpp -o example$(python3-config --extension-suffix)

コンパイルしたC++コードを実行するPythonコードは以下の通りです。

import numpy as np
import example

a = np.asarray([1, 2, 3, 4]) 
b = np.asarray([2, 3, 4, 5])

print(example.add_arrays(a, b)) # [3. 5. 7. 9.]

期待通りの1次元配列 (numpy.ndarray) が返ってきています（戻り値のコピーが行われているかどうかは調べてないです）。

2次元・3次元の場合も期待通りの結果になります。

import numpy as np
import example

a = np.asarray([[1, 2], [3, 4]])
b = np.asarray([[2, 3], [4, 5]])
print(example.add_arrays(a, b))
# [[3. 5.]
#  [7. 9.]]

a = np.asarray([[[1], [2]], [[3], [4]]]) 
b = np.asarray([[[2], [3]], [[4], [5]]]) 
print(example.add_arrays(a, b))
# [[[3.]
#   [5.]]
# 
#  [[7.]
#   [9.]]]

EigenのときはVectorとMatrixを区別する必要がありましたが、xtensorは次元数に関わらず同じ関数が使えるのでよりNumPyに近い使用感を実現できています。 他の関数の使用感は確認していませんが、チートシートを見る限りではかなりNumPyを再現できているように見えます。また、BLASバインディングを提供するxtensor-blasというプロジェクトもあるようなので、速度を求める人の需要もある程度は満たせそうな気がします。もう少し試してみてもいいかもしれません。