diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix
index 310c26911e5e..0cd673918890 100644
--- a/pkgs/development/python-modules/distrax/default.nix
+++ b/pkgs/development/python-modules/distrax/default.nix
@@ -36,6 +36,9 @@ buildPythonPackage rec {
     flit-core
   ];
 
+  pythonRemoveDeps = [
+    "tfp-nightly"
+  ];
   dependencies = [
     absl-py
     chex
diff --git a/pkgs/development/python-modules/rlax/default.nix b/pkgs/development/python-modules/rlax/default.nix
index f4830da2ccd5..8e106b949c65 100644
--- a/pkgs/development/python-modules/rlax/default.nix
+++ b/pkgs/development/python-modules/rlax/default.nix
@@ -2,10 +2,9 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
-  fetchpatch,
 
   # build-system
-  setuptools,
+  flit-core,
 
   # dependencies
   absl-py,
@@ -15,7 +14,6 @@
   jax,
   jaxlib,
   numpy,
-  tensorflow-probability,
 
   # tests
   dm-haiku,
@@ -37,7 +35,7 @@ buildPythonPackage rec {
   };
 
   build-system = [
-    setuptools
+    flit-core
   ];
 
   dependencies = [