mirror of
https://github.com/CHN-beta/nixpkgs.git
synced 2026-01-21 07:21:40 +08:00
33afbf39f6
checkInputs used to be added to nativeBuildInputs. Now we have nativeCheckInputs to do that instead. Doing this treewide change allows to keep hashes identical to before the introduction of nativeCheckInputs.
43 lines
1001 B
Nix
43 lines
1001 B
Nix
{ 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 = "260e5ba01f46b10c579a61393e6c7e546aeae93e";
|
|
sha256 = "sha256-BTHy/jNf6LeV+x3GTI9MDBWLK6A5z2Z1TQyBkHMTeuE=";
|
|
};
|
|
|
|
# Wheel requires only `numpy`, but the import needs `jax`.
|
|
propagatedBuildInputs = [
|
|
jax
|
|
];
|
|
|
|
pythonImportsCheck = [
|
|
"jmp"
|
|
];
|
|
|
|
nativeCheckInputs = [
|
|
jaxlib
|
|
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 ];
|
|
};
|
|
}
|