python312Packages.nutpie: 0.13.4 -> 0.14.2 (#387320)

Author: GaetanLepage

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

@@ -0,0 +1,65 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  hatchling,
+
+  # dependencies
+  equinox,
+  jax,
+  jaxtyping,
+  optax,
+  paramax,
+  tqdm,
+
+  # tests
+  beartype,
+  numpyro,
+  pytest-xdist,
+  pytestCheckHook,
+}:
+
+buildPythonPackage rec {
+  pname = "flowjax";
+  version = "17.1.1";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "danielward27";
+    repo = "flowjax";
+    tag = "v${version}";
+    hash = "sha256-CLtkO7Lr+FpC/RHnqDSmevpJ3/3EfuILrQSbcdNTqsI=";
+  };
+
+  build-system = [
+    hatchling
+  ];
+
+  dependencies = [
+    equinox
+    jax
+    jaxtyping
+    optax
+    paramax
+    tqdm
+  ];
+
+  pythonImportsCheck = [ "flowjax" ];
+
+  nativeCheckInputs = [
+    beartype
+    numpyro
+    pytest-xdist
+    pytestCheckHook
+  ];
+
+  meta = {
+    description = "Distributions, bijections and normalizing flows using Equinox and JAX";
+    homepage = "https://github.com/danielward27/flowjax";
+    changelog = "https://github.com/danielward27/flowjax/releases/tag/v${version}";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}

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

@@ -16,9 +16,12 @@
 
   # tests
   # bridgestan, (not packaged)
+  equinox,
+  flowjax,
   jax,
   jaxlib,
   numba,
+  pytest-timeout,
   pymc,
   pytestCheckHook,
   setuptools,
@@ -27,20 +30,20 @@
 
 buildPythonPackage rec {
   pname = "nutpie";
-  version = "0.13.4";
+  version = "0.14.2";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "pymc-devs";
     repo = "nutpie";
     tag = "v${version}";
-    hash = "sha256-BpKt/EWBefCQUGDxyqF6Xjrj/HUvY4M26gk79R/CyZo=";
+    hash = "sha256-9sHs2JbzVRvAJEoLcz5NxkbElbXblDzxA6oCBtb4yFE=";
   };
 
   cargoDeps = rustPlatform.fetchCargoVendor {
     inherit src;
     name = "${pname}-${version}";
-    hash = "sha256-d0Qk01YwosHlOy3yb/2PV5Op1Wz+yB5ROVbUWfHooEk=";
+    hash = "sha256-j7Vasy4BwOYzH43mWdbu+QsNCdRfvJC6ZvYU8XB5s4E=";
   };
 
   build-system = [
@@ -66,10 +69,13 @@ buildPythonPackage rec {
 
   nativeCheckInputs = [
     # bridgestan
+    equinox
+    flowjax
     numba
     jax
     jaxlib
     pymc
+    pytest-timeout
     pytestCheckHook
     setuptools
     writableTmpDirAsHomeHook

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

@@ -0,0 +1,55 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  hatchling,
+
+  # dependencies
+  equinox,
+  jax,
+  jaxtyping,
+
+  # tests
+  beartype,
+  pytestCheckHook,
+}:
+
+buildPythonPackage rec {
+  pname = "paramax";
+  version = "0.0.0";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "danielward27";
+    repo = "paramax";
+    tag = "v${version}";
+    hash = "sha256-w6F9XuQEwRfOei6gDAyCHt2HUY7I4H92AlEv1Xddv54=";
+  };
+
+  build-system = [
+    hatchling
+  ];
+
+  dependencies = [
+    equinox
+    jax
+    jaxtyping
+  ];
+
+  pythonImportsCheck = [ "paramax" ];
+
+  nativeCheckInputs = [
+    beartype
+    pytestCheckHook
+  ];
+
+  meta = {
+    description = "A small library of paramaterizations and parameter constraints for PyTrees";
+    homepage = "https://github.com/danielward27/paramax";
+    changelog = "https://github.com/danielward27/paramax/releases/tag/v${version}";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -4934,6 +4934,8 @@ self: super: with self; {
 
   flower = callPackage ../development/python-modules/flower { };
 
+  flowjax = callPackage ../development/python-modules/flowjax { };
+
   flowlogs-reader = callPackage ../development/python-modules/flowlogs-reader { };
 
   flowmc = callPackage ../development/python-modules/flowmc { };
@@ -10243,6 +10245,8 @@ self: super: with self; {
 
   param = callPackage ../development/python-modules/param { };
 
+  paramax = callPackage ../development/python-modules/paramax { };
+
   parameter-expansion-patched = callPackage ../development/python-modules/parameter-expansion-patched { };
 
   parameterized = callPackage ../development/python-modules/parameterized { };