diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix
index e7b6257d141d..310c26911e5e 100644
--- a/pkgs/development/python-modules/distrax/default.nix
+++ b/pkgs/development/python-modules/distrax/default.nix
@@ -2,11 +2,19 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
-  fetchpatch2,
+
+  # build-system
+  flit-core,
+
+  # dependencies
+  absl-py,
   chex,
+  jax,
   jaxlib,
   numpy,
   tensorflow-probability,
+
+  # tests
   dm-haiku,
   pytest-xdist,
   pytestCheckHook,
@@ -14,42 +22,24 @@
 
 buildPythonPackage rec {
   pname = "distrax";
-  version = "0.1.5";
+  version = "0.1.7";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "google-deepmind";
     repo = "distrax";
     tag = "v${version}";
-    hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0=";
+    hash = "sha256-R6rGGNzup3O6eZ2z4vygYWTjroE/Irt3aog8Op+0hco=";
   };
 
-  patches = [
-    # TODO: remove at the next release (already on master)
-    (fetchpatch2 {
-      name = "fix-jax-0.6.0-compat";
-      url = "https://github.com/google-deepmind/distrax/commit/c02708ac46518fac00ab2945311e0f2ee32c672c.patch";
-      hash = "sha256-hFNXKoA1b5I6dzhwTRXp/SnkHv89GI6tYwlnBBHwG78=";
-    })
-    # https://github.com/google-deepmind/distrax/pull/289
-    (fetchpatch2 {
-      name = "fix-jax-0.7.0-compat";
-      url = "https://github.com/google-deepmind/distrax/commit/7fc5bd7efff4a7144d175199159f115c3e68a3cf.patch";
-      hash = "sha256-TiD72YIb6ajpaCO1yOGl/+JCuaikQ879Zcpaf2wzMq4=";
-    })
+  build-system = [
+    flit-core
   ];
 
-  # TODO: remove at the next release (already on master)
-  # https://github.com/google-deepmind/distrax/pull/293
-  postPatch = ''
-    substituteInPlace distrax/_src/utils/transformations.py \
-      --replace-fail \
-        "jax.experimental.pjit.pjit_p" \
-        "jex.core.primitives.jit_p"
-  '';
-
   dependencies = [
+    absl-py
     chex
+    jax
     jaxlib
     numpy
     tensorflow-probability