From 5b320a4c0111975011de2d0f53be8dddab2387a5 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 30 Dec 2025 10:59:55 +0000
Subject: [PATCH 1/2] python3Packages.distrax: remove tfp-nightly dep

---
 pkgs/development/python-modules/distrax/default.nix | 3 +++
 1 file changed, 3 insertions(+)

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

From 958785066e240be90aa3658cb98f59df9b3ad8f7 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 30 Dec 2025 11:00:43 +0000
Subject: [PATCH 2/2] python3Packages.rlax: fix build

---
 pkgs/development/python-modules/rlax/default.nix | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

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 = [