Add Jax to ML image?

[rabernat]

Jax is becoming very important / useful to us. Would it be easy to add it to the ML image?

[scottyhq]

Jax is becoming very important / useful to us. Would it be easy to add it to the ML image?

Sure, what is jax used for? I'm not sure the ML image has been actively used much recently. Previously @jhamman and @jsadler2 helped decide what packages are in there. Note that it currently is also missing many of the geoscience packages that have recently been added to pangeo-notebook. But the environment solving has been complicated in the past, so generally the fewer the packages the better it seems.

In any case, to add jax, just create a PR adding a line here:

pangeo-docker-images/ml-notebook/environment.yml

Line 27 in 7e9a109

- jupyterlab-nvdashboard

[rsignell-usgs]

@rabernat would be curious to know what you are going to use this for?
I'm guessing something related to this:
https://pymc-devs.medium.com/the-future-of-pymc3-or-theano-is-dead-long-live-theano-d8005f8a0e9b

[scottyhq]

closed by #196