Jax environment with poetry2nix

Hello! Does anyone have a working development environment with Jax using poetry2nix? I am trying to get one working but I’ve run into two stumbling blocks.

So far, I wrote a micro shell.nix file with python 3.11 and poetry as build inputs. Then I initialized a poetry project in my folder, and made a little main.py file that just imports jax and a “start” function prints out the version . When I poetry run start I get the following error: ImportError: libstdc++.so.6: cannot open shared object file: No such file or directory. I tried passing zlib as a build input but that doesn’t help.

I then tried writing a short flake.nix to try building the project via poetry2nix, but when I nix build I get an infinite recursion error. I poured over the poetry.lock file but didn’t find any obviously recursive dependencies. This is my flake.nix attempt:

  inputs = {
    nixpkgs.url = "github:nixos/nixpkgs/nixpkgs-unstable";
  };

  outputs = { self, nixpkgs }:
  let
    system = "x86_64-linux";
    pkgs = nixpkgs.legacyPackages.${system};
  in
    {
      packages.${system}.default = pkgs.poetry2nix.mkPoetryApplication {
        projectDir = self;
      };
    };
}

I got things working, and if it might help someone out, I’ll post here how I managed to cludge something together. I ended up not using poetry2nix, and just using numtide’s devshells and pip and venv.

nix flake new -t "github:numtide/devshell" project

Then I patched the LD_LIBRARY_PATH environment variable in the devshell.toml as follows:

packages = [
    "gcc",
    "zlib",
    "zlib.dev",
    "stdenv.cc.cc.lib",

    # Python
    "python311",
    "python311Packages.setuptools",
]

[[env]]
name = "CPPFLAGS"
eval = "-I${DEVSHELL_DIR}/include"

[[env]]
name = "LDFLAGS"
eval = "-L${DEVSHELL_DIR}/lib"

[[env]]
name = "LD_LIBRARY_PATH"
eval = "${DEVSHELL_DIR}/lib:${LD_LIBRARY_PATH}"

In packages I added anything that might be needed by jax in terms of C++. Anyway it’s working! I just installed everything else in a virtual environment with pip.