Wei Ji

Jun 2023 update. Tried running the following for aarch64: ``` cd pangeo-notebook/ conda-lock lock -f environment.yml -f ../pangeo-notebook/environment.yml -f ../base-notebook/environment.yml -p linux-aarch64 ``` and it looks like only 2 packages...

Thanks @ngam for getting those last packages over the finish line for `linux-aarch64`! I just tried locking `pangeo-notebook` and it worked :tada: Pytorch and tensorflow doesn't seem ready yet, but...

The latest version of `flax` on conda-forge is still `0.6.1`, see https://anaconda.org/conda-forge/flax/files. It seems like the package has fallen out of maintenance for a while, but there's some effort to...

OK, looks like `flax=0.7.4` and `flax=0.7.5` is available on conda-forge now. However, there's some dependency conflict with `tensorflow`. I tried with this `environment.yml` file: ``` name: pangeo channels: - conda-forge...

Looks like we might need to upgrade from CUDA 11.8 to 12 to get a newer version of `tensorflow=2.15.0` from conda-forge with `libprotobuf~=4.24` that works with `flax>=0.7.4`, see https://github.com/conda-forge/tensorflow-feedstock/pull/367#issuecomment-1890784430.

Right, looks like we'll need to expedite the upgrade to CUDA 12 then as mentioned at https://github.com/pangeo-data/pangeo-docker-images/issues/489#issuecomment-1911393877. Let me open a PR for that (got some free time today), and...

Ok, not as simple as I thought. I tried running `conda-lock` to create a lockfile with a newer version of `tensorflow=2.15.0` and `flax>=0.8.0` with CUDA 12.0, but it errors with:...

Nope, `flax>=0.7.0` doesn't work either ``` ├─ flax >=0.7.0 is installable and it requires │ └─ jax >=0.4.11 with the potential options │ ├─ jax 0.4.11 would require │ │...

Hi @racheetmatai, thanks for opening this bug report. It seems like you've tested docker images up to tag `2023.11.14`. I'm wondering if any of the newer ones, e.g. [`2024.01.03`](https://github.com/pangeo-data/pangeo-docker-images/releases/tag/2024.01.03) which...

KvikIO loads data into `cupy`, but it should technically be possible to zero-copy `cupy` arrays to JAX, Pytorch, or any array library that implements conversion via dlpack or the `__cuda_array_interface__`...