-
Notifications
You must be signed in to change notification settings - Fork 197
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
nb::ndarray<...>
implementation refactored (#721)
This PR refactors the ``nb::ndarray`` implementation to remove certain redundancies. In particular, there were duplicate code paths to process ``ndarray`` template parameters at compile time and at runtime, which are now merged. Significant edits to the documentation are intended to make nd-array bindings more approachable to newcomers. Finally, the refactor was also an opportunity to realize two usability improvements: 1. The constructor to return new nd-arrays from C++ now considers all template arguments: - **Memory order**: ``c_contig``, ``f_contig``. - **Shape**: ``nb::shape<3, 4, 5>``, etc. - **Device type**: ``nb::device::cpu``, ``nb::device::cuda``, etc. - **Framework**: ``nb::numpy``, ``nb::pytorch``, etc. - **Data type**: ``uint64_t``, ``std::complex<double>``, etc. Previously, only the **framework** and **data type** annotations were taken into account when returning nd-arrays, while all of them were examined when *accepting* arrays during overload resolution. This inconsistency was a repeated source of confusion among users. To give an example, the following now works out of the box without the need to redundantly specify the shape and strides to the ``Array`` constructor below: ```cpp using Array = nb::ndarray<float, nb::numpy, nb::shape<4, 4>, nb::f_contig>; struct Matrix4f { float m[4][4]; Array data() { return Array(m); } }; nb::class_<Matrix4f>(m, "Matrix4f") .def("data", &Matrix4f::data, nb::rv_policy::reference_internal); ``` 2. A new nd-array ``.cast()`` method forces the immediate creation of a Python object with the specified target framework and return value policy, while preserving the type signature in return values. This is useful to return temporaries (e.g. stack-allocated memory) from functions. There are two minor but potentially breaking changes: 1. The ndarray type caster now interprets the ``rv_policy::automatic_reference`` return value policy analogously to the ``rv_policy::automatic``, which means that it references a memory region when the user specifies an ``owner``, and it otherwise copies. This makes it safe to use the ``nb::cast()`` and ``nb::ndarray::cast()`` functions that use this policy as a default. 2. The ``nb::any_contig`` memory order annotation, which previously did nothing, now accepts C- or F-contiguous arrays and rejects non-contiguous ones. In both of these cases, the prior convention seems like it would cause bugs/breakage in practice. If nobody depends on this behavior, it should be OK to fix these without a major version bump. A small change without compatibility implications: the `owner` argument has a default `{}` argument again. I think this is reasonably safe because the `automatic_*` return value policies in nanobind will copy the input array if there isn't an owner. This effectively reverts a change from commit 937a1df.
- Loading branch information
Showing
13 changed files
with
1,163 additions
and
773 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.