@@ -0,0 +1,60 @@
1+{ buildPythonPackage
2+, fetchFromGitHub
3+, keras
4+, lib
5+, matplotlib
6+, msgpack
7+, numpy
8+, optax
9+, pytestCheckHook
10+, tensorflow
11+}:
12+
13+buildPythonPackage rec {
14+ pname = "flax";
15+ version = "0.3.6";
16+
17+ src = fetchFromGitHub {
18+ owner = "google";
19+ repo = pname;
20+ rev = "v${version}";
21+ sha256 = "0zvq0vl88hiwmss49bnm7gdmndr1dfza2bcs1fj88a9r7w9dmlsr";
22+ };
23+
24+ propagatedBuildInputs = [
25+ matplotlib
26+ msgpack
27+ numpy
28+ optax
29+ ];
30+
31+ pythonImportsCheck = [
32+ "flax"
33+ ];
34+
35+ checkInputs = [
36+ keras
37+ pytestCheckHook
38+ tensorflow
39+ ];
40+
41+ disabledTestPaths = [
42+ # Docs test, needs extra deps + we're not interested in it.
43+ "docs/_ext/codediff_test.py"
44+
45+ # The tests in `examples` are not designed to be executed from a single test
46+ # session and thus either have the modules that conflict with each other or
47+ # wrong import paths, depending on how they're invoked. Many tests also have
48+ # dependencies that are not packaged in `nixpkgs` (`clu`, `jgraph`,
49+ # `tensorflow_datasets`, `vocabulary`) so the benefits of trying to run them
50+ # would be limited anyway.
51+ "examples/*"
52+ ];
53+
54+ meta = with lib; {
55+ description = "Neural network library for JAX";
56+ homepage = "https://github.com/google/flax";
57+ license = licenses.asl20;
58+ maintainers = with maintainers; [ ndl ];
59+ };
60+}