Add support for max_version, dl_device, copy kwargs in __dlpack__ to match Array API

[Micky774]

Towards #20200

cf. data-apis/array-api#741, data-apis/array-api#602, https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html

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 max_version=None, dl_device=None, copy=None,

This PR also adds new versioning information to jax/_src/dlpack.py:

DLPACK_VERSION = (0, 8)
MIN_DLPACK_VERSION = (0, 5)

[jakevdp]

Minor comment, but in general it's useful to have meaningful commit messages ("Update" doesn't communicate much about what the change includes).

[Micky774]

@yashk2810 @jakevdp This PR should be ready for review again. I've added tests and updated the documentation of to_dlpack.

[jakevdp]

Tests are failing because jax2tf uses the take_ownership parameter removed in this PR: https://github.com/google/jax/blob/abfbb0ae2b85a6c887caa5d11609b347ee6b8f2d/jax/experimental/jax2tf/call_tf.py#L337

(jax2tf tests are not run in github CI because they are too expensive).
I think a sufficient fix will be removing that parameter from the call site, since it's deprecated and currently does nothing.

[Micky774]

@jakevdp Done -- I've also simplified the tests a bit and refactored DLDeviceType to jax._src.typing