@@ -0,0 +1,36 @@
1+{ buildPythonPackage
2+, einops
3+, fetchFromGitHub
4+, jax
5+, jaxlib
6+, lib
7+}:
8+
9+buildPythonPackage rec {
10+ pname = "augmax";
11+ version = "unstable-2022-02-19";
12+ format = "setuptools";
13+
14+ src = fetchFromGitHub {
15+ owner = "khdlr";
16+ repo = pname;
17+ # augmax does not have releases tagged. See https://github.com/khdlr/augmax/issues/5.
18+ rev = "3e5d85d6921a1e519987d33f226bc13f61e04d04";
19+ sha256 = "046n43v7161w7najzlbi0443q60436xv24nh1mv23yw6psqqhx5i";
20+ };
21+
22+ propagatedBuildInputs = [ einops jax ];
23+
24+ # augmax does not have any tests at the time of writing (2022-02-19), but
25+ # jaxlib is necessary for the pythonImportsCheckPhase.
26+ checkInputs = [ jaxlib ];
27+
28+ pythonImportsCheck = [ "augmax" ];
29+
30+ meta = with lib; {
31+ description = "Efficiently Composable Data Augmentation on the GPU with Jax";
32+ homepage = "https://github.com/khdlr/augmax";
33+ license = licenses.asl20;
34+ maintainers = with maintainers; [ samuela ];
35+ };
36+}