Pin jaxlib to use cuda120 build with cuda-nvcc and update README note

[weiji14]

Adding a explicit pin on jaxlib=*=cuda120* to pick up the cuda120 build version, ~as the conda-lock solver was picking up the cpu build of jaxlib instead of the cuda build in #514.~ See:

https://github.com/pangeo-data/pangeo-docker-images/blob/6d4c2abcc5f145135e1d3213043b406fdf058793/ml-notebook/conda-lock.yml#L4166-L4182

Edit: the cuda120 build is actually picked up automatically now as of the 2024.06.02 tag, but setting the cuda120 pin still to be sure that this doesn't break in the future.

Note that jaxlib-0.4.23-cuda120py* has an explicit runtime dependency on cuda-nvcc since https://github.com/conda-forge/jaxlib-feedstock/pull/241, so this should mean users won't have to install cuda-nvcc explicitly anymore.

The cuda-nvcc workaround in the main README.md file has also been removed, in place of a message recommending users to use ml-notebook>=2024.06.02.

Fixes #438

[pangeo-bot]

/condalock Automatically locking new conda environment, building, and testing images...

[github-actions[bot]]

:point_left: Try on Mybinder.org!

[weiji14]

Pulling in jaxlib with the cuda120 build doesn't work yet, see https://github.com/pangeo-data/pangeo-docker-images/actions/runs/9181646795/job/25248830043#step:4:56:

INFO:conda_lock.conda_lock:Using virtual packages from virtual-packages.yml
Locking dependencies for ['linux-64']...
INFO:conda_lock.conda_solver:linux-64 using specs ['cuda-version >=12.0', 'flax >=0.8.0', 'jax', 'jaxlib >=0.4.23 cuda120*', 'jupyterlab-nvdashboard', 'keras-cv', 'tensorflow >=2.15.0 cuda120*', 'adlfs', 'argopy', 'awscli', 'black', 'boto3', 'bottleneck', 'cartopy', 'cdsapi', 'cfgrib', 'cf_xarray', 'ciso', 'cmocean', 'dask-ml', 'datashader', 'descartes', 'earthaccess', 'eofs', 'erddapy', 'esmpy', 'fastjmd95', 'flox', 'fsspec', 'gcm_filters', 'gcsfs', 'gh', 'gh-scoped-creds', 'geocube', 'geopandas', 'geopy', 'geoviews-core', 'git-lfs', 'gsw', 'h5netcdf', 'h5py', 'holoviews', 'hvplot', 'intake', 'intake-esm', 'intake-geopandas', 'intake-stac', 'intake-xarray', 'ipdb', 'ipykernel', 'ipyleaflet', 'ipytree', 'ipywidgets', 'jupyterlab_code_formatter', 'jupyterlab-git', 'jupyterlab-lsp', 'jupyterlab-myst', 'jupyter-panel-proxy', 'jupyter-resource-usage', 'kerchunk', 'line_profiler', 'lxml', 'lz4', 'matplotlib-base', 'memory_profiler', 'metpy', 'nb_conda_kernels', 'nbstripout', 'nc-time-axis', 'netcdf4', 'numbagg', 'numcodecs', 'numpy', 'numpy_groupies', 'odc-stac', 'pandas', 'panel', 'parcels', 'param', 'pop-tools', 'pyarrow', 'pycamhd', 'pydap', 'pystac', 'pystac-client', 'python-blosc', 'python-gist', 'python-graphviz', 'python-lsp-ruff', 'python-xxhash', 'rasterio', 'rechunker', 'rio-cogeo', 'rioxarray', 'ruff', 's3fs', 'satpy', 'scikit-image', 'scikit-learn', 'scipy', 'seaborn', 'sparse', 'snakeviz', 'stackstac', 'tiledb-py', 'timezonefinder', 'watermark', 'xarray', 'xarrayutils', 'xarray-datatree', 'xarray_leaflet', 'xarray-spatial', 'xbatcher', 'xcape', 'xclim', 'xesmf', 'xgboost', 'xgcm', 'xhistogram', 'xmip', 'xmitgcm', 'xpublish', 'xrft', 'xskillscore', 'xxhash', 'zarr', 'python 3.11.*', 'pangeo-notebook 2024.05.20.*', 'pip']
Failed to parse json, Expecting value: line 1 column 1 (char 0)
Could not lock the environment for platform linux-64
Could not solve for environment specs
The following packages are incompatible
├─ jaxlib >=0.4.23 cuda120* is installable with the potential options
│  ├─ jaxlib 0.4.23 would require
│  │  └─ libabseil >=20240116.1,<20240117.0a0 , which can be installed;
│  ├─ jaxlib 0.4.23 would require
│  │  └─ libabseil >=20240116.2,<20240117.0a0 , which can be installed;
│  └─ jaxlib 0.4.23 would require
│     └─ python_abi 3.12.* *_cp312, which requires
│        └─ python 3.12.* *_cpython, which can be installed;
├─ python 3.11**  is not installable because it conflicts with any installable versions previously reported;
└─ tensorflow >=2.15.0 cuda120* is not installable because it requires
   └─ tensorflow-base [2.15.0 cuda120py310heceb7ac_2|2.15.0 cuda120py310heceb7ac_3|...|2.15.0 cuda120py39hf42b710_3], which requires
      └─ libabseil >=20230802.1,<20230803.0a0 , which conflicts with any installable versions previously reported.
{
    "success": false
}

Need to wait for newer version of tensorflow on conda-forge to use libabseil>=20240116, wait for https://github.com/conda-forge/tensorflow-feedstock/pull/372 or https://github.com/conda-forge/tensorflow-feedstock/pull/385