summaryrefslogtreecommitdiff
path: root/pkgs/development/python-modules/dm-sonnet/default.nix
blob: a113283d55915f209f4a7d89bc52fdc55322a0ce (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
{
  lib,
  buildPythonPackage,
  fetchFromGitHub,

  # build-system
  setuptools,

  # dependencies
  dm-tree,
  etils,
  numpy,
  tabulate,
  wrapt,

  # tests
  click,
  docutils,
  keras,
  pytestCheckHook,
  tensorflow,
  tensorflow-datasets,
  tf-keras,
}:

buildPythonPackage rec {
  pname = "dm-sonnet";
  version = "2.0.2";
  pyproject = true;

  src = fetchFromGitHub {
    owner = "deepmind";
    repo = "sonnet";
    tag = "v${version}";
    hash = "sha256-WkloUbqSyPG3cbLG8ktsjdluACkCbUZ7t6rYWst8rs8=";
  };

  build-system = [
    setuptools
  ];

  dependencies = [
    dm-tree
    etils
    numpy
    tabulate
    wrapt
  ]
  ++ etils.optional-dependencies.epath;

  optional-dependencies = {
    tensorflow = [ tensorflow ];
  };

  nativeCheckInputs = [
    click
    docutils
    keras
    pytestCheckHook
    tensorflow
    tensorflow-datasets
    tf-keras
  ];

  # ImportError: `keras.optimizers.legacy` is not supported in Keras 3
  preCheck = ''
    export TF_USE_LEGACY_KERAS=True
  '';

  enabledTestPaths = [
    # Prevent collecting docs/ext/link_tf_api_test.py which fails with:
    # ModuleNotFoundError: No module named 'docs.ext'
    "sonnet"
  ];

  disabledTests = [
    # AssertionError: 2 != 0 : 2 doctests failed
    "test_doctest_sonnet.functional"

    # AssertionError: Not equal to tolerance
    "testComputationAgainstNumPy1"

    # tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute MatMul as input #1(zero-based) was expected to be a float tensor but is a half tensor [Op:MatMul]
    "testComputationAgainstNumPy0"
    "testComputationAgainstNumPy1"
  ];

  pythonImportsCheck = [ "sonnet" ];

  meta = {
    description = "Library for building neural networks in TensorFlow";
    homepage = "https://github.com/deepmind/sonnet";
    changelog = "https://github.com/google-deepmind/sonnet/releases/tag/v${version}";
    license = lib.licenses.asl20;
    maintainers = with lib.maintainers; [ onny ];
  };
}