如何配置Nix以安装jaxlibWithCuda及CUDA相关非自由包？

临时单次生效方案 #

直接在运行nix develop时追加参数，临时允许非自由包：

nix develop --impure --option allowUnfree true --extra-experimental-features "nix-command flakes"

注：--impure是因为Flake纯模式下会忽略系统级的非自由许可配置，加上该参数才能读取相关设置。

单个Flake项目永久配置 #

修改项目根目录的flake.nix，在导入nixpkgs时明确开启非自由包支持，可选全局放开或精确指定允许的包：

全局允许非自由包 #

{
  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; # 替换为你使用的分支
  };

  outputs = { self, nixpkgs }:
  let
    system = "x86_64-linux"; # 根据你的系统架构调整，如aarch64-linux
    pkgs = import nixpkgs {
      inherit system;
      config.allowUnfree = true; # 核心配置：允许所有非自由包
    };
  in
  {
    devShells.${system}.default = pkgs.mkShell {
      packages = [
        pkgs.jaxlibWithCuda
        pkgs.jax
        # 添加其他需要的依赖
      ];
    };
  };

精确指定允许的非自由包（更安全） #

如果不想全局放开非自由包，可以用allowUnfreePredicate只允许需要的包：

pkgs = import nixpkgs {
  inherit system;
  config.allowUnfreePredicate = pkg: builtins.elem pkg.name [
    "cuda_cudart"
    "jaxlibWithCuda"
    # 可添加其他依赖的非自由包名称
  ];
};

系统级永久配置（所有Nix操作生效） #

编辑/etc/nixos/configuration.nix，添加系统级的非自由许可配置：

nix.settings = {
  allowUnfree = true;
  # 或者用精确控制：
  # allowUnfreePredicate = pkg: builtins.elem pkg.name ["cuda_cudart" "jaxlibWithCuda"];
};

保存后执行：

sudo nixos-rebuild switch

之后运行nix develop无需额外参数即可正常拉取jaxlibWithCuda及相关CUDA非自由包。

注：之前在configuration.nix的配置未生效，是因为Flake默认处于纯模式，不会自动继承系统级的非自由许可设置，需通过上述方法明确配置。

内容的提问来源于stack exchange，提问作者Pierre-olivier Gendraud