From b892aeec2483272640abfeeff8bbc508fe26c3b6 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sun, 12 Jan 2025 19:24:00 +0100
Subject: [PATCH] python312Packages.distrax: disable failing tests

---
 .../python-modules/distrax/default.nix         | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix
index 3a84da55a64d..6cae7b75ee6b 100644
--- a/pkgs/development/python-modules/distrax/default.nix
+++ b/pkgs/development/python-modules/distrax/default.nix
@@ -1,7 +1,6 @@
 {
   lib,
   buildPythonPackage,
-  pythonOlder,
   fetchFromGitHub,
   chex,
   jaxlib,
@@ -17,8 +16,6 @@ buildPythonPackage rec {
   version = "0.1.5";
   pyproject = true;
 
-  disabled = pythonOlder "3.9";
-
   src = fetchFromGitHub {
     owner = "google-deepmind";
     repo = "distrax";
@@ -26,7 +23,7 @@ buildPythonPackage rec {
     hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0=";
   };
 
-  buildInputs = [
+  dependencies = [
     chex
     jaxlib
     numpy
@@ -42,6 +39,12 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "distrax" ];
 
   disabledTests = [
+    # NotImplementedError: Primitive 'square' does not have a registered inverse.
+    "test_against_tfp_bijectors_square"
+    "test_log_dets_square__with_device"
+    "test_log_dets_square__without_device"
+    "test_log_dets_square__without_jit"
+
     # AssertionError on numerical values
     # Reported upstream in https://github.com/google-deepmind/distrax/issues/267
     "test_method_with_input_unnormalized_probs__with_device"
@@ -90,10 +93,11 @@ buildPythonPackage rec {
     "distrax/_src/utils/hmm_test.py"
   ];
 
-  meta = with lib; {
+  meta = {
     description = "Probability distributions in JAX";
     homepage = "https://github.com/deepmind/distrax";
-    license = licenses.asl20;
-    maintainers = with maintainers; [ onny ];
+    changelog = "https://github.com/google-deepmind/distrax/releases/tag/v${version}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ onny ];
   };
 }