diff options
author | Alexander Tsvyashchenko <ndl@endl.ch> | 2022-01-01 22:35:46 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-01-01 13:35:46 -0800 |
commit | 44866de4d11506208d7b0b5e35fe1ffaaa0516a2 (patch) | |
tree | b71946751c269a5e0d1029a5fff192120131b08a /pkgs/development/python-modules/jmp | |
parent | 7a1f9c890834a7c7776bc0dabc0071ca53e8475b (diff) |
python3Packages.jmp: init at unstable-2021-10-03 (#152972)
Diffstat (limited to 'pkgs/development/python-modules/jmp')
-rw-r--r-- | pkgs/development/python-modules/jmp/default.nix | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/pkgs/development/python-modules/jmp/default.nix b/pkgs/development/python-modules/jmp/default.nix new file mode 100644 index 0000000000000..dc096b93ae468 --- /dev/null +++ b/pkgs/development/python-modules/jmp/default.nix @@ -0,0 +1,42 @@ +{ buildPythonPackage +, fetchFromGitHub +, jax +, jaxlib +, lib +, pytestCheckHook +}: + +buildPythonPackage rec { + pname = "jmp"; + # As of 2022-01-01, the latest stable version (0.0.2) fails tests with recent JAX versions, + # IIUC it's fixed in https://github.com/deepmind/jmp/commit/4969392f618d7733b265677143d8c81e44085867 + version = "unstable-2021-10-03"; + + src = fetchFromGitHub { + owner = "deepmind"; + repo = pname; + rev = "4b94370b8de29b79d6f840b09d1990b91c1afddd"; + sha256 = "0hh4cmp93wjyidj48gh07vhx2kjvpwd23xvy79bsjn5qaaf6q4cm"; + }; + + # Wheel requires only `numpy`, but the import needs both `jax` and `jaxlib`. + propagatedBuildInputs = [ + jax + jaxlib + ]; + + pythonImportsCheck = [ + "jmp" + ]; + + checkInputs = [ + pytestCheckHook + ]; + + meta = with lib; { + description = "This library implements support for mixed precision training in JAX."; + homepage = "https://github.com/deepmind/jmp"; + license = licenses.asl20; + maintainers = with maintainers; [ ndl ]; + }; +} |