python3Packages.rlax: fix jax 0.10.0 compatibility

pkgs/development/python-modules/rlax/default.nix

@@ -22,18 +22,25 @@
   pytestCheckHook,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "rlax";
   version = "0.1.8";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "google-deepmind";
     repo = "rlax";
-    tag = "v${version}";
+    tag = "v${finalAttrs.version}";
     hash = "sha256-E/zYFd5bfx58FfA3uR7hzRAIs844QzJA8TZTwmwDByk=";
   };
 
+  # TypeError: clip() got an unexpected keyword argument 'a_min'
+  postPatch = ''
+    substituteInPlace rlax/_src/mpo_ops.py \
+      --replace-fail "a_min=" "min="
+  '';
+
   build-system = [
     flit-core
   ];
@@ -82,8 +89,8 @@ buildPythonPackage rec {
   meta = {
     description = "Library of reinforcement learning building blocks in JAX";
     homepage = "https://github.com/deepmind/rlax";
-    changelog = "https://github.com/google-deepmind/rlax/releases/tag/${src.tag}";
+    changelog = "https://github.com/google-deepmind/rlax/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ onny ];
   };
-}
+})