pyrox.dev
1{ buildPythonPackage
2, einops
3, fetchFromGitHub
4, jax
5, jaxlib
6, lib
7}:
8
9buildPythonPackage 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}