python311Packages.distrax: fix type errors

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

@@ -74,15 +74,26 @@ buildPythonPackage rec {
     "distrax/_src/distributions/transformed_test.py"
     "distrax/_src/distributions/uniform_test.py"
     "distrax/_src/utils/transformations_test.py"
+    # https://github.com/google-deepmind/distrax/pull/270
+    "distrax/_src/distributions/deterministic_test.py"
+    "distrax/_src/distributions/epsilon_greedy_test.py"
+    "distrax/_src/distributions/gamma_test.py"
+    "distrax/_src/distributions/greedy_test.py"
+    "distrax/_src/distributions/gumbel_test.py"
+    "distrax/_src/distributions/logistic_test.py"
+    "distrax/_src/distributions/log_stddev_normal_test.py"
+    "distrax/_src/distributions/mvn_diag_test.py"
+    "distrax/_src/distributions/mvn_full_covariance_test.py"
+    "distrax/_src/distributions/mvn_tri_test.py"
+    "distrax/_src/distributions/one_hot_categorical_test.py"
+    "distrax/_src/distributions/softmax_test.py"
+    "distrax/_src/utils/hmm_test.py"
   ];
 
   meta = with lib; {
     description = "Probability distributions in JAX";
     homepage = "https://github.com/deepmind/distrax";
     license = licenses.asl20;
     maintainers = with maintainers; [ onny ];
-    # Several tests fail with:
-    # AssertionError: [Chex] Assertion assert_type failed: Error in type compatibility check
-    broken = true;
   };
 }