Add array-api copy semantics to DLPackManagedTensorToBuffer #29963
+81
−36
Conversation
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
pearu
changed the title
pearu
force-pushed
the
pearu/from_dlpack-forced-copy
branch
from
604aea5 to
51573a4
Compare
pearu
changed the title
pearu
force-pushed
the
pearu/from_dlpack-forced-copy
branch
3 times, most recently
from
e226d31 to
d2c005f
Compare
pearu
force-pushed
the
pearu/from_dlpack-forced-copy
branch
from
d2c005f to
11ed1a3
Compare
superbobry
approved these changes
Jul 7, 2025
| se = str(e) | ||
| i = se.index("is not aligned to") | ||
| if i > 0: | ||
| raise ValueError( |
There was a problem hiding this comment.
I assume we cannot raise this error directly in C++ because it talks about Python-level arguments?
| _arr = jnp.asarray(_buf) | ||
| if copy and jaxlib_version >= (0, 6, 3): | ||
| # dlpack_managed_tensor_to_buffer implements array-api copy | ||
| # semantics, so resetting copy to avoid an unnecessary copy |
There was a problem hiding this comment.
It might be more clear to say that copy was already handled by dlpack_managed_tensor_to_buffer, I think.
Also, nit: add a "." :)
| // require a view (copy == False) or try create a view (copy == None) | ||
| auto result = device.client()->CreateViewOfDeviceBuffer( | ||
| data, shape, *device.default_memory_space(), on_delete_callback, stream); | ||
| if ((result.status().code() != absl::StatusCode::kInvalidArgument) |
There was a problem hiding this comment.
Nit: drop redundant parens?
| TF_ASSIGN_OR_RETURN(auto* memory_space, device.default_memory_space()); | ||
|
|
||
| // Create a copy. | ||
| auto result = device.client()->BufferFromHostBuffer( |
| bool fallback_to_copy = !copy.has_value() && dlmt->dl_tensor.device.device_type == kDLCPU; | ||
|
|
||
| if (!copy.value_or(false)) { | ||
| // require a view (copy == False) or try create a view (copy == None) |
There was a problem hiding this comment.
Nit: I'd nuke this and other similar comments personally, and instead try to change the code to make the conditions more obvious (e.g. on L233).
| @@ -738,7 +738,7 @@ def buffer_to_dlpack_managed_tensor( | |||
| ) -> Any: ... | |||
| @overload | |||
| def dlpack_managed_tensor_to_buffer( | |||
| tensor: Any, device: Device, stream: int | None | |||
| tensor: Any, device: Device, stream: int | None, copy: bool | None | |||
There was a problem hiding this comment.
Add = None, since copy has a default.
google-ml-butler
bot
added
kokoro:force-run
pull ready
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.
As in the title.
Fixes #29810