-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
WIP attempt to update packages but GPU stopped working
- Loading branch information
tartavull
committed
Aug 20, 2023
1 parent
d0516ba
commit 851b743
Showing
6 changed files
with
62 additions
and
145 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,108 +0,0 @@ | ||
{ lib | ||
, absl-py | ||
, blas | ||
, buildPythonPackage | ||
, etils | ||
, fetchFromGitHub | ||
, lapack | ||
, matplotlib | ||
, numpy | ||
, opt-einsum | ||
, pytestCheckHook | ||
, pytest-xdist | ||
, pythonOlder | ||
, scipy | ||
, typing-extensions | ||
, jaxlib | ||
}: | ||
|
||
let | ||
usingMKL = blas.implementation == "mkl" || lapack.implementation == "mkl"; | ||
in | ||
buildPythonPackage rec { | ||
pname = "jax"; | ||
version = "0.4.10"; | ||
format = "setuptools"; | ||
|
||
disabled = pythonOlder "3.7"; | ||
|
||
src = fetchFromGitHub { | ||
owner = "google"; | ||
repo = pname; | ||
rev = "jax-v${version}"; | ||
hash = "sha256-USdEVEcZ28YHDJQDzWzpWdBQQimi27xe5Quc9dESoXw="; | ||
}; | ||
|
||
# jaxlib is _not_ included in propagatedBuildInputs because there are | ||
# different versions of jaxlib depending on the desired target hardware. The | ||
# JAX project ships separate wheels for CPU, GPU, and TPU. Currently only the | ||
# CPU wheel is packaged. | ||
propagatedBuildInputs = [ | ||
absl-py | ||
etils | ||
numpy | ||
opt-einsum | ||
scipy | ||
typing-extensions | ||
jaxlib | ||
] ++ etils.optional-dependencies.epath; | ||
|
||
checkInputs = [ | ||
jaxlib | ||
matplotlib | ||
pytestCheckHook | ||
pytest-xdist | ||
]; | ||
|
||
# high parallelism will result in the tests getting stuck | ||
dontUsePytestXdist = true; | ||
|
||
# NOTE: Don't run the tests in the expiremental directory as they require flax | ||
# which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2. | ||
# Not a big deal, this is how the JAX docs suggest running the test suite | ||
# anyhow. | ||
pytestFlagsArray = [ | ||
"--numprocesses=4" | ||
"-W ignore::DeprecationWarning" | ||
"tests/" | ||
]; | ||
|
||
disabledTests = [ | ||
# Exceeds tolerance when the machine is busy | ||
"test_custom_linear_solve_aux" | ||
] ++ lib.optionals usingMKL [ | ||
# See | ||
# * https://github.com/google/jax/issues/9705 | ||
# * https://discourse.nixos.org/t/getting-different-results-for-the-same-build-on-two-equally-configured-machines/17921 | ||
# * https://github.com/NixOS/nixpkgs/issues/161960 | ||
"test_custom_linear_solve_cholesky" | ||
"test_custom_root_with_aux" | ||
"testEigvalsGrad_shape" | ||
]; | ||
|
||
doCheck = false; # Disable running checks during the build process | ||
|
||
|
||
# See https://github.com/google/jax/issues/11722. This is a temporary fix in | ||
# order to unblock etils, and upgrading jax/jaxlib to the latest version. See | ||
# https://github.com/NixOS/nixpkgs/issues/183173#issuecomment-1204074993. | ||
disabledTestPaths = [ | ||
"tests/api_test.py" | ||
"tests/core_test.py" | ||
"tests/lax_numpy_indexing_test.py" | ||
"tests/lax_numpy_test.py" | ||
"tests/nn_test.py" | ||
"tests/random_test.py" | ||
"tests/sparse_test.py" | ||
]; | ||
|
||
# As of 0.3.22, `import jax` does not work without jaxlib being installed. | ||
pythonImportsCheck = [ ]; | ||
|
||
meta = with lib; { | ||
description = "Differentiate, compile, and transform Numpy code"; | ||
homepage = "https://github.com/google/jax"; | ||
license = licenses.asl20; | ||
maintainers = with maintainers; [ samuela ]; | ||
}; | ||
} | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters