Try to install latest TensorFlow version using mamba

images/tensorflow-notebook/Dockerfile

@@ -17,10 +17,16 @@ RUN mamba install --yes \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 
-# Install tensorflow with pip, on x86_64 tensorflow-cpu
-RUN [[ $(uname -m) = x86_64 ]] && TF_POSTFIX="-cpu" || TF_POSTFIX="" && \
-    pip install --no-cache-dir \
-    "tensorflow${TF_POSTFIX}" && \
+# Can't use `conda-forge` for aarch64 tensorflow:
+# https://github.com/conda-forge/tensorflow-feedstock/issues/136
+RUN if [ "$(uname -m)" = "x86_64" ]; then \
+        mamba install --yes \
+            'tensorflow-cpu' && \
+            mamba clean --all -f -y ; \
+    else \
+        pip install --no-cache-dir \
+            'tensorflow'; \
+    fi && \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 

images/tensorflow-notebook/cuda/Dockerfile

@@ -17,9 +17,12 @@ RUN mamba install --yes \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 
-# Install TensorFlow, CUDA and cuDNN with pip
-RUN pip install --no-cache-dir \
-    "tensorflow[and-cuda]<=2.17.1" && \
+# Install TensorFlow, CUDA and cuDNN
+# Specifying `CONDA_OVERRIDE_CUDA` to use gpu version on a machine with cpu only:
+# https://github.com/conda-forge/tensorflow-feedstock/issues/174
+RUN CONDA_OVERRIDE_CUDA=12.6 mamba install --yes \
+    'tensorflow-gpu' && \
+    mamba clean --all -f -y && \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 

tests/by_image/docker-stacks-foundation/test_packages.py

@@ -39,6 +39,8 @@
     "pytables": "tables",
     "scikit-image": "skimage",
     "scikit-learn": "sklearn",
+    "tensorflow-cpu": "tensorflow",
+    "tensorflow-gpu": "tensorflow",
     # R
     "randomforest": "randomForest",
     "rcurl": "RCurl",