diff --git a/nix/pysisyphus.nix b/nix/pysisyphus.nix
index 72b50f7ff..40463704d 100644
--- a/nix/pysisyphus.nix
+++ b/nix/pysisyphus.nix
@@ -31,6 +31,8 @@
 , pyscf
 , numba
 , optuna
+, jax
+, jaxlib
   # Runtime dependencies
 , runtimeShell
 , jmol
@@ -146,6 +148,8 @@ buildPythonPackage rec {
     pyscf
     numba
     optuna
+    jax
+    jaxlib
   ] # Syscalls
   ++ lib.optional enableXtb xtb
   ++ lib.optional enableTblite tblite
@@ -193,7 +197,6 @@ buildPythonPackage rec {
         "test_ci_opt[RFOptimizer-opt_kwargs1--78.2487951]" # Broken as of PySCF >= 2.3 as a DFT functional definition was changed
         "test_composite_oniom[lib:subst_effect/toluene_minus_H_b3lypG_631g.xyz-2-high_inds1--270.824806805671]" # Slight numerical deviations
         "test_spectrum_ctnum"
-        "test_constrain_h2o_ofix"
       ];
     in
     [