Merge pull request #168325 from samuela/samuela/haiku

python3Packages.dm-haiku: fix build
This commit is contained in:
Samuel Ainsworth 2022-04-13 10:21:28 -07:00 committed by GitHub
commit c57ceea79d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 10 additions and 6 deletions

View file

@ -12,16 +12,14 @@
buildPythonPackage rec { buildPythonPackage rec {
pname = "chex"; pname = "chex";
# As of 2021-12-29, the latest official version has broken tests with jax 0.2.26: version = "0.1.2";
# `AttributeError: module 'jax.interpreters.xla' has no attribute 'xb'`
version = "unstable-2021-12-16";
format = "setuptools"; format = "setuptools";
src = fetchFromGitHub { src = fetchFromGitHub {
owner = "deepmind"; owner = "deepmind";
repo = pname; repo = pname;
rev = "5adc10e0b4218f8ec775567fca38b68bbad42a3a"; rev = "v${version}";
sha256 = "00xib6zv9pwid2q7wcr109qj3fa3g3b852skz8444kw7r0qxy7z3"; hash = "sha256-NtZYOHByKBcKmhRaNULwaQqxfoPRmgbtJ3cFHNfy4E8=";
}; };
propagatedBuildInputs = [ propagatedBuildInputs = [

View file

@ -21,7 +21,7 @@ buildPythonPackage rec {
owner = "deepmind"; owner = "deepmind";
repo = pname; repo = pname;
rev = "v${version}"; rev = "v${version}";
sha256 = "sha256-qvKMeGPiWXvvyV+GZdTWdsC6Wp08AmP8nDtWk7sZtqM="; hash = "sha256-qvKMeGPiWXvvyV+GZdTWdsC6Wp08AmP8nDtWk7sZtqM=";
}; };
propagatedBuildInputs = [ propagatedBuildInputs = [
@ -32,6 +32,7 @@ buildPythonPackage rec {
checkInputs = [ checkInputs = [
chex chex
cloudpickle cloudpickle
dill
dm-tree dm-tree
jaxlib jaxlib
pytest-xdist pytest-xdist
@ -55,6 +56,11 @@ buildPythonPackage rec {
"haiku/_src/integration/jax2tf_test.py" "haiku/_src/integration/jax2tf_test.py"
]; ];
disabledTests = [
# See https://github.com/deepmind/dm-haiku/issues/366.
"test_jit_Recurrent"
];
meta = with lib; { meta = with lib; {
description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet."; description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet.";
homepage = "https://github.com/deepmind/dm-haiku"; homepage = "https://github.com/deepmind/dm-haiku";