Merge remote-tracking branch 'upstream/main' into dev-why

dev_env_builder.nix

@@ -0,0 +1,44 @@
+{ nixpkgs
+, system
+, cudaSupport
+}:
+let
+  pkgs = import nixpkgs {
+    inherit system;
+    config.allowUnfree = true;
+    config.cudaSupport = cudaSupport;
+  };
+  python = pkgs.python310;
+
+  common-dependencies = with python.pkgs; [
+    jax
+    jaxlib
+    optax
+    pyarrow
+  ];
+
+  doc-dependencies = with python.pkgs; [
+    myst-parser
+    numpydoc
+    pydata-sphinx-theme
+    sphinx
+    sphinx-copybutton
+    sphinx-design
+  ];
+
+  test-dependencies = with python.pkgs; [
+    pytest
+    chex
+    gym
+    ray
+    torchvision
+  ];
+in
+pkgs.mkShell {
+  buildInputs = [
+    (python.withPackages
+      (ps: common-dependencies
+        ++ test-dependencies
+        ++ doc-dependencies))
+  ];
+}

fhs.nix

@@ -19,11 +19,11 @@ in
     stdenv
     gcc
 
-    python39
-    python39Packages.pip
-    python39Packages.setuptools
-    python39Packages.wheel
-    python39Packages.virtualenv
+    python310
+    python310Packages.pip
+    python310Packages.setuptools
+    python310Packages.wheel
+    python310Packages.virtualenv
 
     cudaPackages.cudatoolkit
     cudaPackages.cudnn
@@ -32,6 +32,4 @@ in
     SDL2
     glfw
   ];
-
-  runScript = "fish";
 }).env

flake.nix

@@ -1,7 +1,7 @@
 {
   description = "evox";
   inputs = {
-    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+    nixpkgs.url = "github:ogoid/nixpkgs/fixup-missing-cudnn";
     utils.url = "github:numtide/flake-utils";
     flake-compat = {
       url = "github:edolstra/flake-compat";
@@ -14,15 +14,17 @@
     with nixpkgs.lib;
     eachSystem (with system; [ x86_64-linux ]) (system:
       let
-        generic-builder = import ./generic.nix;
-        with-cuda = generic-builder { inherit system nixpkgs; cudaSupport = true; };
-        cpu-only = generic-builder { inherit system nixpkgs; cudaSupport = false; };
-        total = recursiveUpdate with-cuda cpu-only;
+        builder = import ./dev_env_builder.nix;
+        cuda-env = builder { inherit system nixpkgs; cudaSupport = true; };
+        cpu-env = builder { inherit system nixpkgs; cudaSupport = false; };
       in
-      recursiveUpdate total {
-        devShells.default = total.devShells.cpu;
-        packages.default = total.packages.cpu;
-        devShells.fhs = import ./fhs.nix { inherit nixpkgs system; cudaSupport=true; };
+      {
+        devShells.default = cpu-env;
+        devShells.cpu = cpu-env;
+        devShells.cuda = cuda-env;
+        devShells.fhs = import ./fhs.nix {
+          inherit nixpkgs system; cudaSupport=true;
+        };
       }
     );
 }

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
   "jax >= 0.3.0",
   "jaxlib >= 0.3.0",
   "optax >= 0.1.0",
+  "pyarrow >= 10.0.0",
 ]
 
 [project.optional-dependencies]

requirements/docs-requirements.txt

@@ -2,6 +2,7 @@
 jax
 jaxlib
 optax
+pyarrow
 # for docs generation
 pydata-sphinx-theme
 sphinx < 6 # https://github.com/pydata/pydata-sphinx-theme/issues/1094

src/evox/algorithms/mo/__init__.py

@@ -2,3 +2,4 @@
 from .rvea import RVEA
 from .moead import MOEAD
 from .ibea import IBEA
+from .nsga3 import NSGA3