pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 setuptools,
8
9 # dependencies
10 fsspec,
11 lightning-utilities,
12 numpy,
13 packaging,
14 pyyaml,
15 torch,
16 torchmetrics,
17 tqdm,
18 traitlets,
19
20 # tests
21 psutil,
22 pytestCheckHook,
23}:
24
25buildPythonPackage rec {
26 pname = "pytorch-lightning";
27 version = "2.5.5";
28 pyproject = true;
29
30 src = fetchFromGitHub {
31 owner = "Lightning-AI";
32 repo = "pytorch-lightning";
33 tag = version;
34 hash = "sha256-8CDVvgaxnFWO4Fl5lW/+cn/1WZCgVXYys86iOVNYUfY=";
35 };
36
37 preConfigure = ''
38 export PACKAGE_NAME=pytorch
39 '';
40
41 build-system = [ setuptools ];
42
43 dependencies = [
44 fsspec
45 lightning-utilities
46 numpy
47 packaging
48 pyyaml
49 torch
50 torchmetrics
51 tqdm
52 traitlets
53 ]
54 ++ fsspec.optional-dependencies.http;
55
56 nativeCheckInputs = [
57 psutil
58 pytestCheckHook
59 ];
60
61 # Some packages are not in NixPkgs; other tests try to build distributed
62 # models, which doesn't work in the sandbox.
63 doCheck = false;
64
65 pythonImportsCheck = [ "pytorch_lightning" ];
66
67 meta = {
68 description = "Lightweight PyTorch wrapper for machine learning researchers";
69 homepage = "https://github.com/Lightning-AI/pytorch-lightning";
70 changelog = "https://github.com/Lightning-AI/pytorch-lightning/releases/tag/${src.tag}";
71 license = lib.licenses.asl20;
72 maintainers = with lib.maintainers; [ tbenst ];
73 };
74}