diff --git a/pkgs/development/python-modules/ax-platform/default.nix b/pkgs/development/python-modules/ax-platform/default.nix
index a82550a8de9c..aee136f5c4e8 100644
--- a/pkgs/development/python-modules/ax-platform/default.nix
+++ b/pkgs/development/python-modules/ax-platform/default.nix
@@ -22,6 +22,8 @@
   sympy,
 
   # tests
+  jax,
+  numpyro,
   pyfakefs,
   pytestCheckHook,
   sqlalchemy,
@@ -30,14 +32,15 @@
 
 buildPythonPackage (finalAttrs: {
   pname = "ax-platform";
-  version = "1.2.4";
+  version = "1.3.1";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "facebook";
     repo = "ax";
     tag = finalAttrs.version;
-    hash = "sha256-ZeYoLOPF2D1bk89V5/WO/v0UCtDisbAOOL/Su9sB2Fg=";
+    hash = "sha256-8Q/p74rzMAciNUjaTJtf6duslc3TXT7MIjcHt2sffqI=";
   };
 
   env.ALLOW_BOTORCH_LATEST = "1";
@@ -63,6 +66,8 @@ buildPythonPackage (finalAttrs: {
   ++ botorch.optional-dependencies.pymoo;
 
   nativeCheckInputs = [
+    jax
+    numpyro
     pyfakefs
     pytestCheckHook
     sqlalchemy