diff options
Diffstat (limited to 'pkgs/development/python-modules/jaxlib/default.nix')
| -rw-r--r-- | pkgs/development/python-modules/jaxlib/default.nix | 21 |
1 files changed, 11 insertions, 10 deletions
diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix index 7410400ed05a5..de7fea0a26e8a 100644 --- a/pkgs/development/python-modules/jaxlib/default.nix +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -4,7 +4,7 @@ stdenv, # Build-time dependencies: - addOpenGLRunpath, + addDriverRunpath, autoAddDriverRunpath, bazel_6, binutils, @@ -67,16 +67,17 @@ let effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else inputs.stdenv; meta = with lib; { - description = "JAX is Autograd and XLA, brought together for high-performance machine learning research"; + description = "Source-built JAX backend. JAX is Autograd and XLA, brought together for high-performance machine learning research"; homepage = "https://github.com/google/jax"; license = licenses.asl20; maintainers = with maintainers; [ ndl ]; - platforms = platforms.unix; + + # Make this platforms.unix once Darwin is supported. + # The top-level jaxlib now falls back to jaxlib-bin on unsupported platforms. # aarch64-darwin is broken because of https://github.com/bazelbuild/rules_cc/pull/136 # however even with that fix applied, it doesn't work for everyone: # https://github.com/NixOS/nixpkgs/pull/184395#issuecomment-1207287129 - # NOTE: We always build with NCCL; if it is unsupported, then our build is broken. - broken = effectiveStdenv.isDarwin || nccl.meta.unsupported; + platforms = platforms.linux; }; # Bazel wants a merged cudnn at configuration time @@ -239,7 +240,7 @@ let wheel build which - ] ++ lib.optionals effectiveStdenv.isDarwin [ cctools ]; + ] ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [ cctools ]; buildInputs = [ @@ -258,8 +259,8 @@ let snappy zlib ] - ++ lib.optionals effectiveStdenv.isDarwin [ IOKit ] - ++ lib.optionals (!effectiveStdenv.isDarwin) [ nsync ]; + ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [ IOKit ] + ++ lib.optionals (!effectiveStdenv.hostPlatform.isDarwin) [ nsync ]; # We don't want to be quite so picky regarding bazel version postPatch = '' @@ -405,7 +406,7 @@ let TF_SYSTEM_LIBS = lib.concatStringsSep "," ( tf_system_libs - ++ lib.optionals (!effectiveStdenv.isDarwin) [ + ++ lib.optionals (!effectiveStdenv.hostPlatform.isDarwin) [ "nsync" # fails to build on darwin ] ); @@ -413,7 +414,7 @@ let # Note: we cannot do most of this patching at `patch` phase as the deps # are not available yet. Framework search paths aren't added by bintools # hook. See https://github.com/NixOS/nixpkgs/pull/41914. - preBuild = lib.optionalString effectiveStdenv.isDarwin '' + preBuild = lib.optionalString effectiveStdenv.hostPlatform.isDarwin '' export NIX_LDFLAGS+=" -F${IOKit}/Library/Frameworks" substituteInPlace ../output/external/rules_cc/cc/private/toolchain/osx_cc_wrapper.sh.tpl \ --replace "/usr/bin/install_name_tool" "${cctools}/bin/install_name_tool" |