Support Metal DLPack ndarray imports and byte offsets#1338
Open
XXXXRT666 wants to merge 4 commits into
Open
Conversation
hpkfft
reviewed
May 8, 2026
1a56c25 to
e38bc86
Compare
4 tasks
Author
|
Added explicit DLPack byte_offset export support: |
hpkfft
reviewed
May 10, 2026
hpkfft
reviewed
May 10, 2026
hpkfft
reviewed
May 10, 2026
hpkfft
reviewed
May 10, 2026
hpkfft
reviewed
May 10, 2026
hpkfft
reviewed
May 10, 2026
Author
|
Addressed the latest review |
Contributor
|
Thanks. I think it's helpful for module developers that this PR now uses the name
so it's unnecessarily limiting to measure strides in terms of element size. |
4 tasks
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This adds test coverage for importing Metal-backed DLPack arrays into
nb::ndarray, including PyTorch MPS tensors.It also exposes raw DLPack storage metadata for device-backed arrays:
data_handle()data_offset()nb::ndarrayexports can now setDLTensor::byte_offsetexplicitly via a new trailingdata_offsetconstructor argument. This lets producers represent views into opaque device allocations, such as Metalid<MTLBuffer>handles, without copying when the logical array starts after the beginning of the allocation.Tests
needs_torch_mps.byte_offset.