diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix
index 83a4a16d650e..bae7ef5ffb93 100644
--- a/pkgs/development/python-modules/distrax/default.nix
+++ b/pkgs/development/python-modules/distrax/default.nix
@@ -74,6 +74,20 @@ buildPythonPackage rec {
     "distrax/_src/distributions/transformed_test.py"
     "distrax/_src/distributions/uniform_test.py"
     "distrax/_src/utils/transformations_test.py"
+    # https://github.com/google-deepmind/distrax/pull/270
+    "distrax/_src/distributions/deterministic_test.py"
+    "distrax/_src/distributions/epsilon_greedy_test.py"
+    "distrax/_src/distributions/gamma_test.py"
+    "distrax/_src/distributions/greedy_test.py"
+    "distrax/_src/distributions/gumbel_test.py"
+    "distrax/_src/distributions/logistic_test.py"
+    "distrax/_src/distributions/log_stddev_normal_test.py"
+    "distrax/_src/distributions/mvn_diag_test.py"
+    "distrax/_src/distributions/mvn_full_covariance_test.py"
+    "distrax/_src/distributions/mvn_tri_test.py"
+    "distrax/_src/distributions/one_hot_categorical_test.py"
+    "distrax/_src/distributions/softmax_test.py"
+    "distrax/_src/utils/hmm_test.py"
   ];
 
   meta = with lib; {
@@ -81,8 +95,5 @@ buildPythonPackage rec {
     homepage = "https://github.com/deepmind/distrax";
     license = licenses.asl20;
     maintainers = with maintainers; [ onny ];
-    # Several tests fail with:
-    # AssertionError: [Chex] Assertion assert_type failed: Error in type compatibility check
-    broken = true;
   };
 }