Merge pull request #279901 from GaetanLepage/dm-haiku (46ec6ef0) · Commits · nix / nixpkgs

pkgs/development/python-modules/bsuite/default.nix

+9 −1

Original line number	Diff line number	Diff line
		{ lib
		, fetchPypi
		, buildPythonPackage
		, fetchpatch
		, frozendict
		, termcolor
		, matplotlib
		@@ -35,7 +36,14 @@ let bsuite = buildPythonPackage rec {
		hash = "sha256-ak9McvXl7Nz5toUaPaRaJek9lurxiQiIW209GnZEjX0=";
		};

		buildInputs = [
		patches = [
		(fetchpatch { # Convert np.int -> np.int32 since np.int is deprecated (https://github.com/google-deepmind/bsuite/pull/48)
		url = "https://github.com/google-deepmind/bsuite/pull/48/commits/f8d81b2f1c27ef2c8c71ae286001ed879ea306ab.patch";
		hash = "sha256-FXtvVS+U8brulq8Z27+yWIimB+kigGiUOIv1SHb1TA8=";
		})
		];

		propagatedBuildInputs = [
		absl-py
		dm-env
		dm-tree

pkgs/development/python-modules/distrax/default.nix

+2 −0

Original line number	Diff line number	Diff line
		@@ -7,6 +7,7 @@
		, numpy
		, tensorflow-probability
		, dm-haiku
		, pytest-xdist
		, pytestCheckHook
		}:

		@@ -33,6 +34,7 @@ buildPythonPackage rec {

		nativeCheckInputs = [
		dm-haiku
		pytest-xdist
		pytestCheckHook
		];

pkgs/development/python-modules/dm-haiku/default.nix

+73 −28

Original line number	Diff line number	Diff line
		{ buildPythonPackage
		{ lib
		, buildPythonPackage
		, fetchFromGitHub
		, fetchpatch
		, callPackage
		, lib
		, absl-py
		, flax
		, jaxlib
		, jmp
		, numpy
		, tabulate
		, jaxlib
		, pytest-xdist
		, pytestCheckHook
		, bsuite
		, chex
		, cloudpickle
		, dill
		, dm-env
		, dm-tree
		, optax
		, rlax
		, tensorflow
		}:

		buildPythonPackage rec {
		let dm-haiku = buildPythonPackage rec {
		pname = "dm-haiku";
		version = "0.0.10";
		version = "0.0.11";
		format = "setuptools";

		src = fetchFromGitHub {
		owner = "deepmind";
		repo = pname;
		repo = "dm-haiku";
		rev = "refs/tags/v${version}";
		hash = "sha256-EZx3o6PgTeFjTwI9Ko9H39EqPSE0yLWWpsdqX6ALlo4=";
		hash = "sha256-xve1vNsVOC6/HVtzmzswM/Sk3uUNaTtqNAKheFb/tmI=";
		};

		patches = [
		# https://github.com/deepmind/dm-haiku/issues/717
		(fetchpatch {
		name = "remove-typing-extensions.patch";
		url = "https://github.com/deepmind/dm-haiku/commit/c22867db1a3314a382bd2ce36511e2b756dc32a8.patch";
		hash = "sha256-SxJc8FrImwMqTJ5OuJ1f4T+HfHgW/sGqXeIqlxEatlE=";
		})
		# https://github.com/deepmind/dm-haiku/pull/672
		(fetchpatch {
		name = "fix-find-namespace-packages.patch";
		@@ -35,14 +42,12 @@ buildPythonPackage rec {
		})
		];

		outputs = [
		"out"
		"testsout"
		];

		propagatedBuildInputs = [
		absl-py
		flax
		jaxlib
		jmp
		numpy
		tabulate
		];

		@@ -50,17 +55,56 @@ buildPythonPackage rec {
		"haiku"
		];

		postInstall = ''
		mkdir $testsout
		cp -R examples $testsout/examples
		'';
		nativeCheckInputs = [
		bsuite
		chex
		cloudpickle
		dill
		dm-env
		dm-haiku
		dm-tree
		jaxlib
		optax
		pytest-xdist
		pytestCheckHook
		rlax
		tensorflow
		];

		disabledTests = [
		# See https://github.com/deepmind/dm-haiku/issues/366.
		"test_jit_Recurrent"

		# Assertion errors
		"testShapeChecking0"
		"testShapeChecking1"

		# This test requires a more recent version of tensorflow. The current one (2.13) is not enough.
		"test_reshape_convert"

		# This test requires JAX support for double precision (64bit), but enabling this causes several
		# other tests to fail.
		# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
		"test_doctest_haiku.experimental"
		];

		disabledTestPaths = [
		# Those tests requires a more recent version of tensorflow. The current one (2.13) is not enough.
		"haiku/_src/integration/jax2tf_test.py"
		];

		# check in passthru.tests.pytest to escape infinite recursion with bsuite
		doCheck = false;

		passthru.tests = {
		pytest = callPackage ./tests.nix { };
		};
		# check in passthru.tests.pytest to escape infinite recursion with bsuite
		passthru.tests.pytest = dm-haiku.overridePythonAttrs (_: {
		pname = "${pname}-tests";
		doCheck = true;

		# We don't have to install because the only purpose
		# of this passthru test is to, well, test.
		# This fixes having to set `catchConflicts` to false.
		dontInstall = true;
		});

		meta = with lib; {
		description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet.";
		@@ -68,4 +112,5 @@ buildPythonPackage rec {
		license = licenses.asl20;
		maintainers = with maintainers; [ ndl ];
		};
		}
		};
		in dm-haiku

pkgs/development/python-modules/dm-haiku/tests.nix

deleted100644 → 0

+0 −68

Original line number	Diff line number	Diff line
		{ buildPythonPackage
		, dm-haiku
		, chex
		, cloudpickle
		, dill
		, dm-tree
		, jaxlib
		, pytest-xdist
		, pytestCheckHook
		, tensorflow
		, bsuite
		, frozendict
		, dm-env
		, scikit-image
		, rlax
		, distrax
		, tensorflow-probability
		, optax
		}:

		buildPythonPackage {
		pname = "dm-haiku-tests";
		inherit (dm-haiku) version;

		src = dm-haiku.testsout;

		dontBuild = true;
		dontInstall = true;

		nativeCheckInputs = [
		bsuite
		chex
		cloudpickle
		dill
		distrax
		dm-env
		dm-haiku
		dm-tree
		frozendict
		jaxlib
		pytest-xdist
		pytestCheckHook
		optax
		rlax
		scikit-image
		tensorflow
		tensorflow-probability
		];

		disabledTests = [
		# See https://github.com/deepmind/dm-haiku/issues/366.
		"test_jit_Recurrent"
		# Assertion errors
		"test_connect_conv_padding_function_same0"
		"test_connect_conv_padding_function_valid0"
		"test_connect_conv_padding_function_same1"
		"test_connect_conv_padding_function_same2"
		"test_connect_conv_padding_function_valid1"
		"test_connect_conv_padding_function_valid2"
		"test_invalid_axis_ListString"
		"test_invalid_axis_String"
		"test_simple_case"
		"test_simple_case_with_scale"
		"test_slice_axis"
		"test_zero_inputs"
		];

		}

pkgs/development/python-modules/rlax/default.nix

+2 −0

Original line number	Diff line number	Diff line
		@@ -12,6 +12,7 @@
		, tensorflow-probability
		, dm-haiku
		, optax
		, pytest-xdist
		, pytestCheckHook
		}:

		@@ -49,6 +50,7 @@ buildPythonPackage rec {
		nativeCheckInputs = [
		dm-haiku
		optax
		pytest-xdist
		pytestCheckHook
		];

Admin message