diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix
index 80dc392c603d..616dbae7a4fd 100644
--- a/pkgs/development/python-modules/distrax/default.nix
+++ b/pkgs/development/python-modules/distrax/default.nix
@@ -1,21 +1,27 @@
 { lib
-, fetchPypi
 , buildPythonPackage
+, pythonOlder
+, fetchFromGitHub
+, chex
+, jaxlib
 , numpy
 , tensorflow-probability
-, chex
 , dm-haiku
 , pytestCheckHook
-, jaxlib
 }:
 
 buildPythonPackage rec {
   pname = "distrax";
-  version = "0.1.4";
+  version = "0.1.5";
+  pyproject = true;
 
-  src = fetchPypi {
-    inherit pname version;
-    hash = "sha256-klXT5wfnWUGMrf5sQhYqz7Foc/Ou5y4GIFgtTff1ZFQ=";
+  disabled = pythonOlder "3.9";
+
+  src = fetchFromGitHub {
+    owner = "google-deepmind";
+    repo = "distrax";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0=";
   };
 
   buildInputs = [
@@ -34,6 +40,26 @@ buildPythonPackage rec {
     "distrax"
   ];
 
+  disabledTests = [
+    # AssertionError on numerical values
+    # Reported upstream in https://github.com/google-deepmind/distrax/issues/267
+    "test_method_with_input_unnormalized_probs__with_device"
+    "test_method_with_input_unnormalized_probs__with_jit"
+    "test_method_with_input_unnormalized_probs__without_device"
+    "test_method_with_input_unnormalized_probs__without_jit"
+    "test_method_with_value_1d"
+    "test_nested_distributions__with_device"
+    "test_nested_distributions__without_device"
+    "test_nested_distributions__with_jit"
+    "test_nested_distributions__without_jit"
+    "test_stability__with_device"
+    "test_stability__with_jit"
+    "test_stability__without_device"
+    "test_stability__without_jit"
+    "test_von_mises_sample_gradient"
+    "test_von_mises_sample_moments"
+  ];
+
   disabledTestPaths = [
     # TypeErrors
     "distrax/_src/bijectors/tfp_compatible_bijector_test.py"
@@ -54,7 +80,5 @@ buildPythonPackage rec {
     homepage = "https://github.com/deepmind/distrax";
     license = licenses.asl20;
     maintainers = with maintainers; [ onny ];
-    # Broken on all platforms (starting 2022-07-27)
-    broken = true;
   };
 }