From cac71724a8dd608d43e6d5cb126dbdf4db92ee50 Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Fri, 9 Aug 2024 19:40:05 -0400 Subject: [PATCH 1/6] [feature/nanobind] First rough test of nanobind --- nanobind_test/CMakeLists.txt | 38 +++++++++++ nanobind_test/README.md | 13 ++++ nanobind_test/cmake/CPM.cmake | 30 +++++++++ nanobind_test/mdspan_wrapper.cpp | 107 +++++++++++++++++++++++++++++++ nanobind_test/test.py | 53 +++++++++++++++ 5 files changed, 241 insertions(+) create mode 100644 nanobind_test/CMakeLists.txt create mode 100644 nanobind_test/README.md create mode 100644 nanobind_test/cmake/CPM.cmake create mode 100644 nanobind_test/mdspan_wrapper.cpp create mode 100644 nanobind_test/test.py diff --git a/nanobind_test/CMakeLists.txt b/nanobind_test/CMakeLists.txt new file mode 100644 index 0000000..65f8fb5 --- /dev/null +++ b/nanobind_test/CMakeLists.txt @@ -0,0 +1,38 @@ +cmake_minimum_required(VERSION 3.27) + +set(CMAKE_EXPORT_COMPILE_COMMANDS + TRUE + CACHE BOOL "Export compile commands to build directory" FORCE) + +include(cmake/CPM.cmake) + +# Set C++ standard +set(CMAKE_CXX_STANDARD 23) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) +# Need to enforce -fPIC across whole project to build shared libraries +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +find_package( + Python 3.12 + COMPONENTS Interpreter Development + REQUIRED) + +# TODO different compilation flags for Debug/Release modes; also pull in +# target_compile_options for fast_pauli target defined below +if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE + Release + CACHE STRING "Choose the type of build." FORCE) + set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" + "MinSizeRel" "RelWithDebInfo") +endif() + +# Dependencies +cpmaddpackage("gh:wjakob/nanobind#v2.0.0") +cpmaddpackage("gh:fmtlib/fmt#10.2.1") +cpmaddpackage("gh:kokkos/mdspan#b885a2c60ad42f9e1aaa0d317a38105b950cbed0") + +# add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${nanobind_BUILD_DIR}) +nanobind_add_module(mdspan_wrapper mdspan_wrapper.cpp) +target_link_libraries(mdspan_wrapper PRIVATE mdspan fmt::fmt) diff --git a/nanobind_test/README.md b/nanobind_test/README.md new file mode 100644 index 0000000..60238c2 --- /dev/null +++ b/nanobind_test/README.md @@ -0,0 +1,13 @@ +# `nb::ndarray` :handshake: `std::mdspan` + +How to connect `nanobind`'s `nb::ndarray` and `std::mdspan`. + +Properties of `nb::ndarray` +- work with numpy, torch, tensorflow, jax, cupy, and anything that supports DLPack +- Supports zero-copy exchange using two protocols: 1) buffer protocol and 2) DLPack +- Can return data that is owned by a Python object + +## Useful guides from `nanobind` +- https://github.com/wjakob/nanobind/blob/b1531b9397c448c6b784520a4f052608f28a5e8d/include/nanobind/eigen/dense.h#L93C7-L93C24 +- https://github.com/wjakob/nanobind/blob/b1531b9397c448c6b784520a4f052608f28a5e8d/include/nanobind/eigen/dense.h#L124 +- https://github.com/wjakob/nanobind/blob/b1531b9397c448c6b784520a4f052608f28a5e8d/include/nanobind/eigen/dense.h#L242-L245 \ No newline at end of file diff --git a/nanobind_test/cmake/CPM.cmake b/nanobind_test/cmake/CPM.cmake new file mode 100644 index 0000000..e847861 --- /dev/null +++ b/nanobind_test/cmake/CPM.cmake @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: MIT +# +# SPDX-FileCopyrightText: Copyright (c) 2019-2023 Lars Melchior and contributors + +set(CPM_DOWNLOAD_VERSION 0.39.0) +set(CPM_HASH_SUM + "66639bcac9dd2907b2918de466783554c1334446b9874e90d38e3778d404c2ef") + +if(CPM_SOURCE_CACHE) + set(CPM_DOWNLOAD_LOCATION + "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake") +elseif(DEFINED ENV{CPM_SOURCE_CACHE}) + set(CPM_DOWNLOAD_LOCATION + "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake") +else() + set(CPM_DOWNLOAD_LOCATION + "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake") +endif() + +# Expand relative path. This is important if the provided path contains a tilde +# (~) +get_filename_component(CPM_DOWNLOAD_LOCATION ${CPM_DOWNLOAD_LOCATION} ABSOLUTE) + +file( + DOWNLOAD + https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake + ${CPM_DOWNLOAD_LOCATION} + EXPECTED_HASH SHA256=${CPM_HASH_SUM}) + +include(${CPM_DOWNLOAD_LOCATION}) diff --git a/nanobind_test/mdspan_wrapper.cpp b/nanobind_test/mdspan_wrapper.cpp new file mode 100644 index 0000000..e34b473 --- /dev/null +++ b/nanobind_test/mdspan_wrapper.cpp @@ -0,0 +1,107 @@ +#include +#include +#include + +#include +#include +#include +#include + +namespace nb = nanobind; +using namespace nb::literals; + +template +std::mdspan> +ndarray_cast_from_py(nb::ndarray &a) { + + std::array shape; + for (size_t i = 0; i < ndim; ++i) { + shape[i] = a.shape(i); + } + + return std::mdspan>(a.data(), shape); +} + +void scale_mdspan(std::mdspan> a, + double scale) { + for (size_t i = 0; i < a.extent(0); ++i) { + for (size_t j = 0; j < a.extent(1); ++j) { + for (size_t k = 0; k < a.extent(2); ++k) { + a[i, j, k] *= scale; + } + } + } +} + +void scale_ndarray(nb::ndarray &a, double scale) { + scale_mdspan(ndarray_cast_from_py(a), scale); +} + +template struct PauliOp { + // As n_pauli_strings x n_qubits + std::mdspan> pauli_strings; + std::mdspan> coeffs; + + PauliOp(nb::ndarray &ps, nb::ndarray &c) { + pauli_strings = ndarray_cast_from_py(ps); + coeffs = ndarray_cast_from_py(c); + } + + bool operator==(const PauliOp &other) const { + if (pauli_strings.extent(0) != other.pauli_strings.extent(0) || + pauli_strings.extent(1) != other.pauli_strings.extent(1)) { + return false; + } + + if (coeffs.extent(0) != other.coeffs.extent(0)) { + return false; + } + + fmt::println("Dimensions match"); + + for (size_t i = 0; i < pauli_strings.extent(0); ++i) { + for (size_t j = 0; j < pauli_strings.extent(1); ++j) { + if (pauli_strings[i, j] != other.pauli_strings[i, j]) { + return false; + } + } + } + + fmt::println("Pauli strings match"); + + for (size_t i = 0; i < coeffs.extent(0); ++i) { + if (coeffs[i] != other.coeffs[i]) { + fmt::println("{} != {}", coeffs[i], other.coeffs[i]); + return false; + } + } + + fmt::println("Coeffs match"); + + return true; + } + + void scale(T scale) { + for (size_t i = 0; i < coeffs.extent(0); ++i) { + coeffs[i] *= scale; + } + } + + void print() { + fmt::print("Coeffs["); + for (size_t i = 0; i < coeffs.extent(0); ++i) { + fmt::print("{}, ", coeffs[i]); + } + fmt::print("]\n"); + } +}; + +NB_MODULE(mdspan_wrapper, m) { + // + m.def("scale_ndarray", &scale_ndarray, "a"_a.noconvert(), "scale"_a); + nb::class_>(m, "PauliOp") + .def(nb::init &, nb::ndarray &>()) + .def("scale", &PauliOp::scale, "scale"_a) + .def("print", &PauliOp::print) + .def("__eq__", &PauliOp::operator==); +} \ No newline at end of file diff --git a/nanobind_test/test.py b/nanobind_test/test.py new file mode 100644 index 0000000..bbc3992 --- /dev/null +++ b/nanobind_test/test.py @@ -0,0 +1,53 @@ +import sys +import numpy as np +import pytest + +sys.path.append("build") + +import mdspan_wrapper + + +def test_scale_ndarray() -> None: + size = 10 + scale = 1.34 + + arr_np = np.random.rand(size, size, size) + arr_np_scaled = arr_np * scale + + arr_mdspan = arr_np.copy() + mdspan_wrapper.scale_ndarray(arr_mdspan, scale) + + np.testing.assert_allclose(arr_mdspan, arr_np_scaled) + + +def test_fake_pauli_op() -> None: + pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32) + coeffs = np.array([1.0, 2.0, 3.0]) + op = mdspan_wrapper.PauliOp(pauli_strings, coeffs) + + print(coeffs) + + op.scale(2.0) + print(coeffs) + + +@pytest.fixture +def my_op() -> mdspan_wrapper.PauliOp: + pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32) + coeffs = np.array([1.0, 2.0, 3.0]) + return mdspan_wrapper.PauliOp(pauli_strings, coeffs) + + +def test_fake_pauli_op_scope(my_op) -> None: + my_op.print() + my_op.scale(2.0) + my_op.print() + my_op.print() + my_op.print() + + pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32) + coeffs = np.array([1.0, 2.0, 3.0]) * 2 + op2 = mdspan_wrapper.PauliOp(pauli_strings, coeffs) + op2.print() + + assert my_op == op2 From 22246d9672a2bb8e25fc611780caf25274769814 Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Tue, 20 Aug 2024 13:44:22 -0400 Subject: [PATCH 2/6] [feature/nanobind] Switching to owned data for structure params, tests still not working --- nanobind_test/mdspan_wrapper.cpp | 36 ++++++++++++++---------------- nanobind_test/test.py | 38 ++++++++++++++++++-------------- 2 files changed, 38 insertions(+), 36 deletions(-) diff --git a/nanobind_test/mdspan_wrapper.cpp b/nanobind_test/mdspan_wrapper.cpp index e34b473..6726c1f 100644 --- a/nanobind_test/mdspan_wrapper.cpp +++ b/nanobind_test/mdspan_wrapper.cpp @@ -5,7 +5,9 @@ #include #include #include +#include #include +#include namespace nb = nanobind; using namespace nb::literals; @@ -22,29 +24,28 @@ ndarray_cast_from_py(nb::ndarray &a) { return std::mdspan>(a.data(), shape); } -void scale_mdspan(std::mdspan> a, - double scale) { - for (size_t i = 0; i < a.extent(0); ++i) { - for (size_t j = 0; j < a.extent(1); ++j) { - for (size_t k = 0; k < a.extent(2); ++k) { - a[i, j, k] *= scale; - } - } - } -} - -void scale_ndarray(nb::ndarray &a, double scale) { - scale_mdspan(ndarray_cast_from_py(a), scale); -} - template struct PauliOp { // As n_pauli_strings x n_qubits std::mdspan> pauli_strings; std::mdspan> coeffs; + std::vector _coeffs; PauliOp(nb::ndarray &ps, nb::ndarray &c) { pauli_strings = ndarray_cast_from_py(ps); - coeffs = ndarray_cast_from_py(c); + + // Make space + size_t size = 1; + for (size_t i = 0; i < c.ndim(); ++i) { + size *= c.shape(i); + } + _coeffs.resize(size); + + // Use std::span to get easy iterators to std::copy + std::span _span(c.data(), size); + std::copy(_span.begin(), _span.end(), _coeffs.begin()); + + // Create mdspan on data that this structure owns + coeffs = std::mdspan>(_coeffs.data(), size); } bool operator==(const PauliOp &other) const { @@ -56,7 +57,6 @@ template struct PauliOp { if (coeffs.extent(0) != other.coeffs.extent(0)) { return false; } - fmt::println("Dimensions match"); for (size_t i = 0; i < pauli_strings.extent(0); ++i) { @@ -66,7 +66,6 @@ template struct PauliOp { } } } - fmt::println("Pauli strings match"); for (size_t i = 0; i < coeffs.extent(0); ++i) { @@ -98,7 +97,6 @@ template struct PauliOp { NB_MODULE(mdspan_wrapper, m) { // - m.def("scale_ndarray", &scale_ndarray, "a"_a.noconvert(), "scale"_a); nb::class_>(m, "PauliOp") .def(nb::init &, nb::ndarray &>()) .def("scale", &PauliOp::scale, "scale"_a) diff --git a/nanobind_test/test.py b/nanobind_test/test.py index bbc3992..4d58d82 100644 --- a/nanobind_test/test.py +++ b/nanobind_test/test.py @@ -7,19 +7,6 @@ import mdspan_wrapper -def test_scale_ndarray() -> None: - size = 10 - scale = 1.34 - - arr_np = np.random.rand(size, size, size) - arr_np_scaled = arr_np * scale - - arr_mdspan = arr_np.copy() - mdspan_wrapper.scale_ndarray(arr_mdspan, scale) - - np.testing.assert_allclose(arr_mdspan, arr_np_scaled) - - def test_fake_pauli_op() -> None: pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32) coeffs = np.array([1.0, 2.0, 3.0]) @@ -39,15 +26,32 @@ def my_op() -> mdspan_wrapper.PauliOp: def test_fake_pauli_op_scope(my_op) -> None: + print() my_op.print() my_op.scale(2.0) + print("After scaling") my_op.print() + + op2 = mdspan_wrapper.PauliOp( + np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32), + np.array([1.0, 2.0, 3.0]) * 2, + ) + print("op2") + op2.print() + + my_op.scale(2.0) my_op.print() + + op3 = mdspan_wrapper.PauliOp( + np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32), + np.array([1.0, 2.0, 3.0]) * 4, + ) + print("op3") + op3.print() + + print("my_op") my_op.print() - pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32) - coeffs = np.array([1.0, 2.0, 3.0]) * 2 - op2 = mdspan_wrapper.PauliOp(pauli_strings, coeffs) + print("op2") op2.print() - assert my_op == op2 From 9f34829079c7fa75680fb5079bf70f386efdb023 Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Fri, 23 Aug 2024 15:34:28 -0400 Subject: [PATCH 3/6] [feature/nanobind] Minimal version appears to be working --- nanobind_test/CMakeLists.txt | 1 + nanobind_test/mdspan_wrapper.cpp | 54 +++++++++++++++++++++----------- nanobind_test/test.py | 46 ++++++++++++++++----------- 3 files changed, 64 insertions(+), 37 deletions(-) diff --git a/nanobind_test/CMakeLists.txt b/nanobind_test/CMakeLists.txt index 65f8fb5..fe2064f 100644 --- a/nanobind_test/CMakeLists.txt +++ b/nanobind_test/CMakeLists.txt @@ -12,6 +12,7 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_EXTENSIONS OFF) # Need to enforce -fPIC across whole project to build shared libraries set(CMAKE_POSITION_INDEPENDENT_CODE ON) +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Werror -g3") find_package( Python 3.12 diff --git a/nanobind_test/mdspan_wrapper.cpp b/nanobind_test/mdspan_wrapper.cpp index 6726c1f..55d7e4a 100644 --- a/nanobind_test/mdspan_wrapper.cpp +++ b/nanobind_test/mdspan_wrapper.cpp @@ -2,11 +2,14 @@ #include #include +#include #include #include #include #include #include +#include +#include #include namespace nb = nanobind; @@ -24,6 +27,22 @@ ndarray_cast_from_py(nb::ndarray &a) { return std::mdspan>(a.data(), shape); } +template +std::pair, std::array> +cast_ndarray_to_blob(nb::ndarray &a) { + // Shape info + size_t size = 1; + std::array shape; + for (size_t i = 0; i < a.ndim(); ++i) { + shape[i] = a.shape(i); + size *= a.shape(i); + } + + // Copy the raw data + std::vector _data(size); + std::memcpy(_data.data(), a.data(), size * sizeof(T)); + return std::make_pair(_data, shape); +} template struct PauliOp { // As n_pauli_strings x n_qubits std::mdspan> pauli_strings; @@ -32,20 +51,9 @@ template struct PauliOp { PauliOp(nb::ndarray &ps, nb::ndarray &c) { pauli_strings = ndarray_cast_from_py(ps); - - // Make space - size_t size = 1; - for (size_t i = 0; i < c.ndim(); ++i) { - size *= c.shape(i); - } - _coeffs.resize(size); - - // Use std::span to get easy iterators to std::copy - std::span _span(c.data(), size); - std::copy(_span.begin(), _span.end(), _coeffs.begin()); - - // Create mdspan on data that this structure owns - coeffs = std::mdspan>(_coeffs.data(), size); + std::array shape; + std::tie(_coeffs, shape) = cast_ndarray_to_blob(c); + coeffs = std::mdspan>(_coeffs.data(), shape); } bool operator==(const PauliOp &other) const { @@ -57,7 +65,7 @@ template struct PauliOp { if (coeffs.extent(0) != other.coeffs.extent(0)) { return false; } - fmt::println("Dimensions match"); + // fmt::println("Dimensions match"); for (size_t i = 0; i < pauli_strings.extent(0); ++i) { for (size_t j = 0; j < pauli_strings.extent(1); ++j) { @@ -66,16 +74,15 @@ template struct PauliOp { } } } - fmt::println("Pauli strings match"); + // fmt::println("Pauli strings match"); for (size_t i = 0; i < coeffs.extent(0); ++i) { if (coeffs[i] != other.coeffs[i]) { - fmt::println("{} != {}", coeffs[i], other.coeffs[i]); return false; } } - fmt::println("Coeffs match"); + // fmt::println("Coeffs match"); return true; } @@ -86,6 +93,12 @@ template struct PauliOp { } } + void multiply_coeff(std::mdspan> other) { + for (size_t i = 0; i < coeffs.extent(0); ++i) { + coeffs[i] *= other[i]; + } + } + void print() { fmt::print("Coeffs["); for (size_t i = 0; i < coeffs.extent(0); ++i) { @@ -101,5 +114,8 @@ NB_MODULE(mdspan_wrapper, m) { .def(nb::init &, nb::ndarray &>()) .def("scale", &PauliOp::scale, "scale"_a) .def("print", &PauliOp::print) - .def("__eq__", &PauliOp::operator==); + .def("__eq__", &PauliOp::operator==) + .def("multiply_coeff", [](PauliOp &op, nb::ndarray &c) { + op.multiply_coeff(ndarray_cast_from_py(c)); + }); } \ No newline at end of file diff --git a/nanobind_test/test.py b/nanobind_test/test.py index 4d58d82..de7e143 100644 --- a/nanobind_test/test.py +++ b/nanobind_test/test.py @@ -26,32 +26,42 @@ def my_op() -> mdspan_wrapper.PauliOp: def test_fake_pauli_op_scope(my_op) -> None: - print() - my_op.print() my_op.scale(2.0) - print("After scaling") - my_op.print() - op2 = mdspan_wrapper.PauliOp( np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32), np.array([1.0, 2.0, 3.0]) * 2, ) - print("op2") - op2.print() - my_op.scale(2.0) - my_op.print() + # my_op.scale(2.0) + # print("After scaling, after op2") + # my_op.print() + + # op3 = mdspan_wrapper.PauliOp( + # np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32), + # np.array([1.0, 2.0, 3.0]) * 4, + # ) + # print("op3") + # op3.print() + + # print("my_op") + # my_op.print() + + # print("op2") + # op2.print() + assert my_op == op2 - op3 = mdspan_wrapper.PauliOp( + +def test_multiply_coeff(my_op) -> None: + coeffs = np.array([1.0, 2.0, 3.0]) + my_op.multiply_coeff(coeffs) + op2 = mdspan_wrapper.PauliOp( np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32), - np.array([1.0, 2.0, 3.0]) * 4, + np.array([1.0, 4.0, 9.0]), ) - print("op3") - op3.print() + assert my_op == op2 - print("my_op") - my_op.print() - print("op2") - op2.print() - assert my_op == op2 +def test_multiply_coeff_diff_type(my_op) -> None: + coeffs = np.array([1.0, 2.0, 3.0], dtype=np.complex128) + with pytest.raises(TypeError): + my_op.multiply_coeff(coeffs) From eb60824cf98b3b1711887658da56e9f202d8466b Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Mon, 26 Aug 2024 12:08:19 -0400 Subject: [PATCH 4/6] [feature/nanobind] Minimal version of returning ndarrays working --- nanobind_test/CMakeLists.txt | 2 + nanobind_test/mdspan_wrapper.cpp | 68 ++++++++++++++++++++++++++++++-- nanobind_test/test.py | 19 +++++++++ 3 files changed, 86 insertions(+), 3 deletions(-) diff --git a/nanobind_test/CMakeLists.txt b/nanobind_test/CMakeLists.txt index fe2064f..8c16009 100644 --- a/nanobind_test/CMakeLists.txt +++ b/nanobind_test/CMakeLists.txt @@ -6,6 +6,8 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS include(cmake/CPM.cmake) +project(nanobind_test LANGUAGES CXX) + # Set C++ standard set(CMAKE_CXX_STANDARD 23) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/nanobind_test/mdspan_wrapper.cpp b/nanobind_test/mdspan_wrapper.cpp index 55d7e4a..1f7006d 100644 --- a/nanobind_test/mdspan_wrapper.cpp +++ b/nanobind_test/mdspan_wrapper.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -106,6 +107,12 @@ template struct PauliOp { } fmt::print("]\n"); } + + void return_coeffs(std::mdspan> &out) { + for (size_t i = 0; i < coeffs.extent(0); ++i) { + out[i] = coeffs[i]; + } + } }; NB_MODULE(mdspan_wrapper, m) { @@ -115,7 +122,62 @@ NB_MODULE(mdspan_wrapper, m) { .def("scale", &PauliOp::scale, "scale"_a) .def("print", &PauliOp::print) .def("__eq__", &PauliOp::operator==) - .def("multiply_coeff", [](PauliOp &op, nb::ndarray &c) { - op.multiply_coeff(ndarray_cast_from_py(c)); + .def("multiply_coeff", + [](PauliOp &op, nb::ndarray &c) { + op.multiply_coeff(ndarray_cast_from_py(c)); + }) + .def("return_coeffs", + [](PauliOp &op, nb::ndarray &out) { + auto out_mdspan = ndarray_cast_from_py(out); + op.return_coeffs(out_mdspan); + }) + .def("return_coeffs_owning", [](PauliOp &op) { + struct Temp { + std::vector data; + }; + + Temp *tmp = new Temp{op._coeffs}; + + fmt::println("copied data: [{}]", fmt::join(tmp->data, ", ")); + std::cout << std::flush; + + nb::capsule deleter( + tmp, [](void *data) noexcept { delete static_cast(data); }); + + return nb::ndarray( + /*data*/ tmp->data.data(), + /*shape */ {tmp->data.size()}, + /*deleter*/ deleter); }); -} \ No newline at end of file + + m.def("return_coeffs", + [](size_t n) { + std::vector data(n); + for (size_t i = 0; i < n; ++i) { + data[i] = i; + } + + struct Temp { + std::vector data; + }; + + Temp *tmp = new Temp{data}; + + nb::capsule deleter(tmp, [](void *data) noexcept { + delete static_cast(data); + }); + + return nb::ndarray( + /*data*/ tmp->data.data(), + /*shape */ {tmp->data.size()}, + /*deleter*/ deleter); + + // nb::ndarray out{data.data(), {data.size()}, deleter}; + + // nb::object res = nb::cast(out, nb::rv_policy::copy); + // return res; + } + + /**/ + ); +} diff --git a/nanobind_test/test.py b/nanobind_test/test.py index de7e143..dec05d0 100644 --- a/nanobind_test/test.py +++ b/nanobind_test/test.py @@ -65,3 +65,22 @@ def test_multiply_coeff_diff_type(my_op) -> None: coeffs = np.array([1.0, 2.0, 3.0], dtype=np.complex128) with pytest.raises(TypeError): my_op.multiply_coeff(coeffs) + + +def test_return_coeffs_non_owning(my_op) -> None: + c = np.zeros(3) + print(c) + my_op.return_coeffs(c) + + print(c) + np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0])) + + +def test_return_coeffs_owning(my_op) -> None: + c = my_op.return_coeffs_owning() + print("returned array", c) + np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0])) + + +def test_standalone() -> None: + np.testing.assert_allclose(mdspan_wrapper.return_coeffs(3), np.array([0, 1.0, 2.0])) From 3ae9f5614d4ceedd0032516d341124590fb9bd4f Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Mon, 26 Aug 2024 15:22:51 -0400 Subject: [PATCH 5/6] [feature/nanobind] Excluding nanobind_test dir from mypy --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 2a5a016..5285a5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,7 @@ strict_optional = true ignore_missing_imports = true explicit_package_bases = true -exclude = [] +exclude = ["nanobind_test"] [tool.ruff] exclude = [ From bf9b055124c0dae3fbe15b9b87d9311879cb543e Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Mon, 26 Aug 2024 15:31:05 -0400 Subject: [PATCH 6/6] [feature/nanobind] Making ruff/mypy happy --- nanobind_test/test.py | 11 ++++++----- pyproject.toml | 3 ++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/nanobind_test/test.py b/nanobind_test/test.py index dec05d0..9367982 100644 --- a/nanobind_test/test.py +++ b/nanobind_test/test.py @@ -1,4 +1,5 @@ import sys + import numpy as np import pytest @@ -25,7 +26,7 @@ def my_op() -> mdspan_wrapper.PauliOp: return mdspan_wrapper.PauliOp(pauli_strings, coeffs) -def test_fake_pauli_op_scope(my_op) -> None: +def test_fake_pauli_op_scope(my_op: mdspan_wrapper.PauliOp) -> None: my_op.scale(2.0) op2 = mdspan_wrapper.PauliOp( np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32), @@ -51,7 +52,7 @@ def test_fake_pauli_op_scope(my_op) -> None: assert my_op == op2 -def test_multiply_coeff(my_op) -> None: +def test_multiply_coeff(my_op: mdspan_wrapper.PauliOp) -> None: coeffs = np.array([1.0, 2.0, 3.0]) my_op.multiply_coeff(coeffs) op2 = mdspan_wrapper.PauliOp( @@ -61,13 +62,13 @@ def test_multiply_coeff(my_op) -> None: assert my_op == op2 -def test_multiply_coeff_diff_type(my_op) -> None: +def test_multiply_coeff_diff_type(my_op: mdspan_wrapper.PauliOp) -> None: coeffs = np.array([1.0, 2.0, 3.0], dtype=np.complex128) with pytest.raises(TypeError): my_op.multiply_coeff(coeffs) -def test_return_coeffs_non_owning(my_op) -> None: +def test_return_coeffs_non_owning(my_op: mdspan_wrapper.PauliOp) -> None: c = np.zeros(3) print(c) my_op.return_coeffs(c) @@ -76,7 +77,7 @@ def test_return_coeffs_non_owning(my_op) -> None: np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0])) -def test_return_coeffs_owning(my_op) -> None: +def test_return_coeffs_owning(my_op: mdspan_wrapper.PauliOp) -> None: c = my_op.return_coeffs_owning() print("returned array", c) np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0])) diff --git a/pyproject.toml b/pyproject.toml index 5285a5a..adc7588 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,7 @@ strict_optional = true ignore_missing_imports = true explicit_package_bases = true -exclude = ["nanobind_test"] +exclude = [] [tool.ruff] exclude = [ @@ -94,6 +94,7 @@ exclude = [ "venv", ".venv", "fast_pauli/__version__.py", + "nanobind_test" ] line-length = 88