From e42671544144634b10bad7fd0865405ebe0936c8 Mon Sep 17 00:00:00 2001
From: Jack Betteridge <j.betteridge@imperial.ac.uk>
Date: Mon, 23 Sep 2024 10:44:00 +0100
Subject: [PATCH] Move both ML frameworks from vanilla to firedrake

---
 docker/Dockerfile.firedrake | 7 +++++--
 docker/Dockerfile.vanilla   | 2 --
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile.firedrake b/docker/Dockerfile.firedrake
index 63f6ce8e72..490a16e59f 100644
--- a/docker/Dockerfile.firedrake
+++ b/docker/Dockerfile.firedrake
@@ -11,14 +11,17 @@ WORKDIR /home/firedrake
 RUN bash -c "source firedrake/bin/activate; \
     firedrake-update \
         --documentation-dependencies \
+        --jax \
         --netgen \
         --slepc \
-        --install femlium \
+        --torch \
+        --install asQ \
         --install defcon \
         --install fascd \
+        --install femlium \
         --install gadopt \
         --install gusto \
         --install icepack \
         --install irksome \
-        --install asQ \
         --install thetis"
+
diff --git a/docker/Dockerfile.vanilla b/docker/Dockerfile.vanilla
index bb41d793d7..2de26b1916 100644
--- a/docker/Dockerfile.vanilla
+++ b/docker/Dockerfile.vanilla
@@ -12,8 +12,6 @@ RUN curl -O https://raw.githubusercontent.com/firedrakeproject/firedrake/master/
 RUN bash -c "python3 firedrake-install \
     --no-package-manager \
     --disable-ssh \
-    --torch \
-    --jax \
     --honour-petsc-dir \
     --mpicc=$MPICH_DIR/mpicc \
     --mpicxx=$MPICH_DIR/mpicxx \