google / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
29.75k stars 2.71k forks source link

Add masked arrays #8979

Open donno2048 opened 2 years ago

donno2048 commented 2 years ago

Please: Add masked arrays (numpy.ma.array) it doesn't have to even be a new feature, it can be done by simply adding mask and fill_value arguments to the existing arrays (jax.numpy.array)

mattjj commented 2 years ago

Thanks for the feature request!

I think adapting everything to work with masked arrays might take quite a bit of work. This may remain out of scope for us for a while. There's a lot of other high-priority work on our plates.

Some of this functionality could likely be implemented as a separate library on top of JAX itself.

donno2048 commented 1 year ago

Is this planned for any time in the future?