Skip to content

Combining JAX and xarray for large scale computing #8164

Answered by keewis
Aceticia asked this question in Q&A
Discussion options

You must be logged in to vote

if you repeatedly would be converting to / from jax, it would probably be best to keep everything in jax and have xarray wrap that.

According to jax-ml/jax#17107, support for doing so seems to be close but not quite there yet. In particular we'd need jax to implement the array API (i.e., define __array_namespace__) – see jax-ml/jax#16099 a in-progress PR.

Replies: 6 comments 28 replies

Comment options

You must be logged in to vote
0 replies
Answer selected by Aceticia
Comment options

You must be logged in to vote
1 reply
@keewis
Comment options

Comment options

You must be logged in to vote
5 replies
@keewis
Comment options

@TomNicholas
Comment options

@keewis
Comment options

@TomNicholas
Comment options

@zazbone
Comment options

Comment options

You must be logged in to vote
21 replies
@allen-adastra
Comment options

@TomNicholas
Comment options

@allen-adastra
Comment options

@shoyer
Comment options

@allen-adastra
Comment options

Comment options

You must be logged in to vote
0 replies
Comment options

You must be logged in to vote
1 reply
@keewis
Comment options

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Category
Q&A
Labels
None yet
8 participants