python312Packages.rlax: skip failing tests (#385584)

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

@@ -39,6 +39,9 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "distrax" ];
 
   disabledTests = [
+    # Flaky: AssertionError: 1 not less than 0.7000000000000001
+    "test_von_mises_sample_uniform_ks_test"
+
     # NotImplementedError: Primitive 'square' does not have a registered inverse.
     "test_against_tfp_bijectors_square"
     "test_log_dets_square__with_device"

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

@@ -3,6 +3,11 @@
   buildPythonPackage,
   fetchFromGitHub,
   fetchpatch,
+
+  # build-system
+  setuptools,
+
+  # dependencies
   absl-py,
   chex,
   distrax,
@@ -11,6 +16,8 @@
   jaxlib,
   numpy,
   tensorflow-probability,
+
+  # tests
   dm-haiku,
   optax,
   pytest-xdist,
@@ -20,7 +27,7 @@
 buildPythonPackage rec {
   pname = "rlax";
   version = "0.1.6";
-  format = "setuptools";
+  pyproject = true;
 
   src = fetchFromGitHub {
     owner = "google-deepmind";
@@ -38,7 +45,11 @@ buildPythonPackage rec {
     })
   ];
 
-  propagatedBuildInputs = [
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
     absl-py
     chex
     distrax
@@ -59,7 +70,14 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "rlax" ];
 
   disabledTests = [
-    # RuntimeErrors
+    # AssertionError: Array(2, dtype=int32) != 0
+    "test_categorical_sample__with_device"
+    "test_categorical_sample__with_jit"
+    "test_categorical_sample__without_device"
+    "test_categorical_sample__without_jit"
+
+    # RuntimeError: Attempted to set 4 devices, but 1 CPUs already available:
+    # ensure that `set_n_cpu_devices` is executed before any JAX operation.
     "test_cross_replica_scatter_add0"
     "test_cross_replica_scatter_add1"
     "test_cross_replica_scatter_add2"
@@ -73,10 +91,11 @@ buildPythonPackage rec {
     "test_unnormalize_linear"
   ];
 
-  meta = with lib; {
+  meta = {
     description = "Library of reinforcement learning building blocks in JAX";
     homepage = "https://github.com/deepmind/rlax";
-    license = licenses.asl20;
-    maintainers = with maintainers; [ onny ];
+    changelog = "https://github.com/google-deepmind/rlax/releases/tag/v${version}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ onny ];
   };
 }