Skip to content

Commit bd6203b

Browse files
add specific optional arguments to Array4.__dlpack__
Signed-off-by: Roelof Groenewald <rgroenewald@tae.com>
1 parent 34e1c00 commit bd6203b

File tree

1 file changed

+10
-9
lines changed

1 file changed

+10
-9
lines changed

src/Base/Array4.H

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -231,13 +231,12 @@ namespace pyAMReX
231231
// https://docs.cupy.dev/en/stable/user_guide/interoperability.html#dlpack-data-exchange-protocol
232232
.def("__dlpack__", [](
233233
Array4<T> const &a4,
234-
py::kwargs kwargs
235-
/* TODO: Handle keyword arguments
236-
[[maybe_unused]] py::handle stream,
237-
[[maybe_unused]] std::tuple<int, int> max_version,
238-
[[maybe_unused]] std::tuple<DLDeviceType, int32_t> dl_device,
239-
[[maybe_unused]] bool copy
240-
*/
234+
/* TODO: Handle keyword arguments */
235+
[[maybe_unused]] std::optional<py::handle> stream = std::nullopt,
236+
[[maybe_unused]] std::optional<std::tuple<int, int>> max_version = std::nullopt,
237+
[[maybe_unused]] std::optional<std::tuple<DLDeviceType, int32_t>> dl_device = std::nullopt,
238+
[[maybe_unused]] std::optional<bool> copy = std::nullopt
239+
241240
)
242241
{
243242
// Allocate shape/strides arrays
@@ -281,8 +280,10 @@ namespace pyAMReX
281280
}
282281
);
283282
},
284-
//py::arg("stream") = py::none(),
285-
// ... other args & their defaults
283+
py::arg("stream") = py::none(),
284+
py::arg("max_version") = py::none(),
285+
py::arg("dl_device") = py::none(),
286+
py::arg("copy") = py::none(),
286287
R"doc(
287288
DLPack protocol for zero-copy tensor exchange.
288289
See https://dmlc.github.io/dlpack/latest/ for details.

0 commit comments

Comments
 (0)