1 files changed, 30 insertions, 20 deletions
diff --git a/pkgs/development/python-modules/ax/default.nix b/pkgs/development/python-modules/ax/default.nix
index f0bba36c3c2fc..dddcc6e1ebeb5 100644
--- a/pkgs/development/python-modules/ax/default.nix
+++ b/pkgs/development/python-modules/ax/default.nix
@@ -1,33 +1,36 @@
-{ lib
-, buildPythonPackage
-, fetchFromGitHub
-, botorch
-, ipywidgets
-, jinja2
-, pandas
-, plotly
-, setuptools
-, setuptools-scm
-, typeguard
-, wheel
-, hypothesis
-, mercurial
-, pyfakefs
-, pytestCheckHook
-, yappi
-, pyre-extensions
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  ax,
+  botorch,
+  ipywidgets,
+  jinja2,
+  pandas,
+  plotly,
+  python,
+  setuptools,
+  setuptools-scm,
+  typeguard,
+  wheel,
+  hypothesis,
+  mercurial,
+  pyfakefs,
+  pytestCheckHook,
+  yappi,
+  pyre-extensions,
 }:
 
 buildPythonPackage rec {
   pname = "ax";
-  version = "0.3.7";
+  version = "0.4.0";
   format = "pyproject";
 
   src = fetchFromGitHub {
     owner = "facebook";
     repo = pname;
     rev = "refs/tags/${version}";
-    hash = "sha256-UOX8Q7F5c9J1KSfQUy8aoL5CNb9YmEI4X7Fd5B16UqQ=";
+    hash = "sha256-dj6Gig8N4oLtcZLwPl4QDHG/FwA2nFBtYxSARnWiJJU=";
   };
 
   nativeBuildInputs = [
@@ -73,9 +76,16 @@ buildPythonPackage rec {
     "test_get_standard_plots_moo"
     # AssertionError: Expected 'warning' to be called once. Called 3 times
     "test_validate_kwarg_typing"
+    # uses torch.equal
+    "test_convert_observations"
   ];
   pythonImportsCheck = [ "ax" ];
 
+  # Many portions of the test suite fail under Python 3.12
+  doCheck = lib.versions.majorMinor python.version != "3.12";
+
+  passthru.tests.check = ax.overridePythonAttrs { doCheck = true; };
+
   meta = with lib; {
     description = "Ax is an accessible, general-purpose platform for understanding, managing, deploying, and automating adaptive experiments";
     homepage = "https://ax.dev/";