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