diff --git a/pkgs/development/python-modules/lineax/default.nix b/pkgs/development/python-modules/lineax/default.nix new file mode 100644 index 0000000000000..8f547bc1d9d76 --- /dev/null +++ b/pkgs/development/python-modules/lineax/default.nix @@ -0,0 +1,66 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + hatchling, + + # dependencies + equinox, + jax, + jaxtyping, + typing-extensions, + + # tests + beartype, + pytest, + python, +}: + +buildPythonPackage rec { + pname = "lineax"; + version = "0.0.6"; + pyproject = true; + + src = fetchFromGitHub { + owner = "patrick-kidger"; + repo = "lineax"; + rev = "refs/tags/v${version}"; + hash = "sha256-rM3H+q75F98eEIJkEszWgxD5C5vGK5RlYtVv8GD/VC0="; + }; + + build-system = [ hatchling ]; + + dependencies = [ + equinox + jax + jaxtyping + typing-extensions + ]; + + pythonImportsCheck = [ "lineax" ]; + + nativeCheckInputs = [ + beartype + pytest + ]; + + # Intentionaly not using pytest directly as it leads to JAX out-of-memory'ing + # https://github.com/patrick-kidger/lineax/blob/1909d190c1963d5f2d991508c1b2714f2266048b/tests/README.md + checkPhase = '' + runHook preCheck + + ${python.interpreter} -m tests + + runHook postCheck + ''; + + meta = { + description = "Linear solvers in JAX and Equinox"; + homepage = "https://github.com/patrick-kidger/lineax"; + changelog = "https://github.com/patrick-kidger/lineax/releases/tag/v${version}"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +} diff --git a/pkgs/development/python-modules/optimistix/default.nix b/pkgs/development/python-modules/optimistix/default.nix new file mode 100644 index 0000000000000..a7cb8b7c5e7e6 --- /dev/null +++ b/pkgs/development/python-modules/optimistix/default.nix @@ -0,0 +1,63 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + hatchling, + + # dependencies + equinox, + jax, + jaxtyping, + lineax, + typing-extensions, + + # checks + beartype, + jaxlib, + optax, + pytestCheckHook, + pytest-xdist, +}: + +buildPythonPackage rec { + pname = "optimistix"; + version = "0.0.8"; + pyproject = true; + + src = fetchFromGitHub { + owner = "patrick-kidger"; + repo = "optimistix"; + rev = "refs/tags/v${version}"; + hash = "sha256-0ehNApiBoiAb8LFBW81ZCRPsjTVQG8zRTVSAp7oHF8w="; + }; + + build-system = [ hatchling ]; + + dependencies = [ + equinox + jax + jaxtyping + lineax + typing-extensions + ]; + + pythonImportsCheck = [ "optimistix" ]; + + nativeCheckInputs = [ + beartype + jaxlib + optax + pytestCheckHook + pytest-xdist + ]; + + meta = { + description = "Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox"; + homepage = "https://github.com/patrick-kidger/optimistix"; + changelog = "https://github.com/patrick-kidger/optimistix/releases/tag/v${version}"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 0938e8a343018..6a2dbc777c089 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -7226,6 +7226,8 @@ self: super: with self; { linear-operator = callPackage ../development/python-modules/linear-operator { }; + lineax = callPackage ../development/python-modules/lineax { }; + linecache2 = callPackage ../development/python-modules/linecache2 { }; lineedit = callPackage ../development/python-modules/lineedit { }; @@ -9408,6 +9410,8 @@ self: super: with self; { optax = callPackage ../development/python-modules/optax { }; + optimistix = callPackage ../development/python-modules/optimistix { }; + optimum = callPackage ../development/python-modules/optimum { }; optree = callPackage ../development/python-modules/optree { };