Add support for `device` and `copy` kwargs in `from_dlpack` to match Array API #20175

Micky774 · 2024-03-11T18:24:02Z

Towards #20200

Note

"In principle, arbitrary cross-device copies could be allowed too, but the consensus in data-apis/array-api#626 was that limiting to device-to-host copies is enough for now". This PR includes the optional device-to-device transfer.

Default behavior is preserved when device=None, copy=None

jakevdp

Looks great! A few minor comments

jax/_src/dlpack.py

jax/_src/numpy/lax_numpy.py

jax/experimental/array_api/_creation_functions.py

jax/numpy/__init__.pyi

jakevdp · 2024-03-11T22:23:33Z

Lint errors might indicate a problem:

jax/_src/dlpack.py:130: error: Argument 2 to "dlpack_managed_tensor_to_buffer" has incompatible type "Client"; expected "Device"  [arg-type]
jax/_src/dlpack.py:130: error: Argument 3 to "dlpack_managed_tensor_to_buffer" has incompatible type "Client | None"; expected "int | None"  [arg-type]

Micky774 · 2024-03-11T22:44:26Z

Lint errors might indicate a problem:

jax/_src/dlpack.py:130: error: Argument 2 to "dlpack_managed_tensor_to_buffer" has incompatible type "Client"; expected "Device"  [arg-type]
jax/_src/dlpack.py:130: error: Argument 3 to "dlpack_managed_tensor_to_buffer" has incompatible type "Client | None"; expected "int | None"  [arg-type]

This is a bug in the xla_extension stub annotations. I'll open a PR in XLA to resolve this. For now, I've added an inline ignore.

jakevdp

Looks great!

jakevdp · 2024-03-12T01:22:05Z

Internal pytype tests are failing with many variations of this error:

File "/jax/_src/third_party/scipy/interpolate.py", line 4, in <module>: Couldn't import pyi for 'jax.numpy' [pyi-error]
  No xla_client.Device in module jax._src.lib, referenced from 'jax.numpy'

jakevdp · 2024-03-12T19:06:13Z

Can you change your commit message to something more informative? Thanks!

…o_buffer` in `python/xla_extension` Imported from GitHub PR openxla/xla#10433 Encountered bug in google/jax#20175 (see this [comment](google/jax#20175 (comment))). This adjusts the stub file to properly overload `dlpack_managed_tensor_to_buffer` so that both signatures can be checked against. Copybara import of the project: -- 75cabb5149b4a0bdd9e819fac0ea6a0ba756bff6 by Meekail Zain <[email protected]>: Update Merging this change closes #10433 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10433 from Micky774:type_update 75cabb5149b4a0bdd9e819fac0ea6a0ba756bff6 PiperOrigin-RevId: 615173747

…o_buffer` in `python/xla_extension` Imported from GitHub PR #10433 Encountered bug in google/jax#20175 (see this [comment](google/jax#20175 (comment))). This adjusts the stub file to properly overload `dlpack_managed_tensor_to_buffer` so that both signatures can be checked against. Copybara import of the project: -- 75cabb5 by Meekail Zain <[email protected]>: Update Merging this change closes #10433 COPYBARA_INTEGRATE_REVIEW=#10433 from Micky774:type_update 75cabb5 PiperOrigin-RevId: 615973838

…o_buffer` in `python/xla_extension` Imported from GitHub PR openxla/xla#10433 Encountered bug in google/jax#20175 (see this [comment](google/jax#20175 (comment))). This adjusts the stub file to properly overload `dlpack_managed_tensor_to_buffer` so that both signatures can be checked against. Copybara import of the project: -- 75cabb5149b4a0bdd9e819fac0ea6a0ba756bff6 by Meekail Zain <[email protected]>: Update Merging this change closes #10433 PiperOrigin-RevId: 615973838

Micky774 · 2024-03-18T12:21:51Z

Internal pytype tests are failing with many variations of this error:

File "/jax/_src/third_party/scipy/interpolate.py", line 4, in <module>: Couldn't import pyi for 'jax.numpy' [pyi-error]
  No xla_client.Device in module jax._src.lib, referenced from 'jax.numpy'

@jakevdp are the internal tests still failing? If so, I will update the typing to use the _Device type to avoid letting this get stalled.

jakevdp · 2024-03-18T16:04:15Z

Still failing:

File "third_party/py/jax/__init__.py", line 163, in <module>: Couldn't import pyi for 'jax.numpy' [pyi-error]
  Can't find pyi for 'jaxlib.xla_client', referenced from 'jax.numpy'

I think the issue is that trying to depend on xla_client doesn't have an interface file. It's why we use Device = Any in other similar locations, e.g. here:

jax/jax/numpy/__init__.pyi

Lines 21 to 22 in aaeeaf5

    
           # TODO(jakevdp): use xla_client.Device here 
        
           _Device = Any

and here:

jax/jax/_src/basearray.pyi

Lines 22 to 23 in aaeeaf5

    
           # TODO: alias this to xla_client.Traceback 
        
           Device = Any

jakevdp

Looks good - could we add some test coverage for the new arguments in array_interoperability_test.py?

tests/array_interoperability_test.py

jakevdp · 2024-04-04T17:18:29Z

This looks good, and is probably ready to merge more-or-less. However there's a subtlety here that I've been thinking about with @yashk2810 – the issue is that the behavior of device=XXX under jit and other transformations is kind of ambiguous: currently device_put is a no-op in this context, which means that this function will silently ignore the device. I think in the short-term we'd prefer to make that an error: from_dlpack should basically fail within any transformation, becuase its semantics are impure: it's reading an external buffer that's not tracked by JAX's normal tracing mechanisms, so e.g. if the buffer changes between the first and second call to the function, cacheing semantics may lead to incorrect outputs.

All of this is somewhat second-order though, so we should probably merge this change and iterate from there.

tests/array_interoperability_test.py

Micky774 force-pushed the array_api branch from 24f219a to 506d520 Compare March 11, 2024 18:25

jakevdp requested changes Mar 11, 2024

View reviewed changes

Micky774 force-pushed the array_api branch 2 times, most recently from cd8ec29 to 40cb7d8 Compare March 11, 2024 22:15

Micky774 force-pushed the array_api branch from 40cb7d8 to d71ac98 Compare March 11, 2024 22:44

Micky774 mentioned this pull request Mar 11, 2024

Fixes type annotation overload of dlpack_managed_tensor_to_buffer in python/xla_extension openxla/xla#10433

Closed

jakevdp approved these changes Mar 11, 2024

View reviewed changes

google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Mar 11, 2024

kokoro-team removed the kokoro:force-run label Mar 11, 2024

Micky774 mentioned this pull request Mar 12, 2024

Update jax.experimental.array_api to v2023.12 API #20200

Open

13 tasks

Micky774 force-pushed the array_api branch from d71ac98 to 2583757 Compare March 12, 2024 18:40

copybara-service bot mentioned this pull request Mar 12, 2024

PR #10433: Fixes type annotation overload of dlpack_managed_tensor_to_buffer in python/xla_extension tensorflow/tensorflow#63555

Merged

Micky774 force-pushed the array_api branch from 0170b35 to 3da84bc Compare March 18, 2024 18:01

jakevdp reviewed Mar 18, 2024

View reviewed changes

jakevdp self-assigned this Mar 18, 2024

Micky774 force-pushed the array_api branch from 3da84bc to 0ce7883 Compare March 22, 2024 19:07

jakevdp reviewed Apr 1, 2024

View reviewed changes

tests/array_interoperability_test.py Outdated Show resolved Hide resolved

Micky774 force-pushed the array_api branch from 0ce7883 to 7e7044a Compare April 3, 2024 19:12

Micky774 force-pushed the array_api branch from 7e7044a to 03338d0 Compare April 3, 2024 20:15

jakevdp approved these changes Apr 4, 2024

View reviewed changes

google-ml-butler bot added the kokoro:force-run label Apr 4, 2024

kokoro-team removed the kokoro:force-run label Apr 4, 2024

jakevdp requested changes Apr 4, 2024

View reviewed changes

tests/array_interoperability_test.py Outdated Show resolved Hide resolved

Update from_dlpack to match array API 2023

2b1c3de

Micky774 force-pushed the array_api branch from 03338d0 to 2b1c3de Compare April 4, 2024 22:52

jakevdp approved these changes Apr 5, 2024

View reviewed changes

google-ml-butler bot added the kokoro:force-run label Apr 5, 2024

kokoro-team removed the kokoro:force-run label Apr 5, 2024

copybara-service bot merged commit f37e503 into google:main Apr 5, 2024
12 of 14 checks passed

Micky774 deleted the array_api branch April 7, 2024 23:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for `device` and `copy` kwargs in `from_dlpack` to match Array API #20175

Add support for `device` and `copy` kwargs in `from_dlpack` to match Array API #20175

Micky774 commented Mar 11, 2024 •

edited

Loading

jakevdp left a comment

jakevdp commented Mar 11, 2024

Micky774 commented Mar 11, 2024

jakevdp left a comment

jakevdp commented Mar 12, 2024

jakevdp commented Mar 12, 2024

Micky774 commented Mar 18, 2024

jakevdp commented Mar 18, 2024 •

edited

Loading

jakevdp left a comment

jakevdp commented Apr 4, 2024

Add support for device and copy kwargs in from_dlpack to match Array API #20175

Add support for device and copy kwargs in from_dlpack to match Array API #20175

Conversation

Micky774 commented Mar 11, 2024 • edited Loading

jakevdp left a comment

Choose a reason for hiding this comment

jakevdp commented Mar 11, 2024

Micky774 commented Mar 11, 2024

jakevdp left a comment

Choose a reason for hiding this comment

jakevdp commented Mar 12, 2024

jakevdp commented Mar 12, 2024

Micky774 commented Mar 18, 2024

jakevdp commented Mar 18, 2024 • edited Loading

jakevdp left a comment

Choose a reason for hiding this comment

jakevdp commented Apr 4, 2024

Add support for `device` and `copy` kwargs in `from_dlpack` to match Array API #20175

Add support for `device` and `copy` kwargs in `from_dlpack` to match Array API #20175

Micky774 commented Mar 11, 2024 •

edited

Loading

jakevdp commented Mar 18, 2024 •

edited

Loading