pyrox.dev
1{ buildPythonPackage
2, einops
3, fetchFromGitHub
4, jax
5, jaxlib
6, lib
7, pythonOlder
8, setuptools
9}:
10
11buildPythonPackage rec {
12 pname = "augmax";
13 version = "0.3.1";
14 pyproject = true;
15
16 disbaled = pythonOlder "3.6";
17
18 src = fetchFromGitHub {
19 owner = "khdlr";
20 repo = "augmax";
21 rev = "refs/tags/v${version}";
22 hash = "sha256-hDNNoE7KVaH3jrXZA4A8f0UoQJPl6KHA3XwMc3Ccohk=";
23 };
24
25 nativeBuildInputs = [
26 setuptools
27 ];
28
29 propagatedBuildInputs = [ einops jax ];
30
31 # augmax does not have any tests at the time of writing (2022-02-19), but
32 # jaxlib is necessary for the pythonImportsCheckPhase.
33 nativeCheckInputs = [ jaxlib ];
34
35 pythonImportsCheck = [ "augmax" ];
36
37 meta = with lib; {
38 description = "Efficiently Composable Data Augmentation on the GPU with Jax";
39 homepage = "https://github.com/khdlr/augmax";
40 changelog = "https://github.com/khdlr/augmax/blob/${src.rev}/CHANGELOG.md";
41 license = licenses.asl20;
42 maintainers = with maintainers; [ samuela ];
43 };
44}