pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 setuptools,
6 wheel,
7 gpuctypes,
8 numpy,
9 tqdm,
10 hypothesis,
11 librosa,
12 onnx,
13 pillow,
14 pytest-xdist,
15 pytestCheckHook,
16 safetensors,
17 sentencepiece,
18 tiktoken,
19 torch,
20 transformers,
21}:
22
23buildPythonPackage rec {
24 pname = "tinygrad";
25 version = "0.8.0";
26 pyproject = true;
27
28 src = fetchFromGitHub {
29 owner = "tinygrad";
30 repo = "tinygrad";
31 rev = "refs/tags/v${version}";
32 hash = "sha256-QAccZ79qUbe27yUykIf22WdkxYUlOffnMlShakKfp60=";
33 };
34
35 nativeBuildInputs = [
36 setuptools
37 wheel
38 ];
39
40 propagatedBuildInputs = [
41 gpuctypes
42 numpy
43 tqdm
44 ];
45
46 pythonImportsCheck = [ "tinygrad" ];
47
48 nativeCheckInputs = [
49 hypothesis
50 librosa
51 onnx
52 pillow
53 pytest-xdist
54 pytestCheckHook
55 safetensors
56 sentencepiece
57 tiktoken
58 torch
59 transformers
60 ];
61
62 preCheck = ''
63 export HOME=$(mktemp -d)
64 '';
65
66 disabledTests = [
67 # Require internet access
68 "test_benchmark_openpilot_model"
69 "test_bn_alone"
70 "test_bn_linear"
71 "test_bn_mnist"
72 "test_car"
73 "test_chicken"
74 "test_chicken_bigbatch"
75 "test_conv_mnist"
76 "testCopySHMtoDefault"
77 "test_data_parallel_resnet"
78 "test_e2e_big"
79 "test_fetch_small"
80 "test_huggingface_enet_safetensors"
81 "test_linear_mnist"
82 "test_load_convnext"
83 "test_load_enet"
84 "test_load_enet_alt"
85 "test_load_llama2bfloat"
86 "test_load_resnet"
87 "test_openpilot_model"
88 "test_resnet"
89 "test_shufflenet"
90 "test_transcribe_batch12"
91 "test_transcribe_batch21"
92 "test_transcribe_file1"
93 "test_transcribe_file2"
94 "test_transcribe_long"
95 "test_transcribe_long_no_batch"
96 "test_vgg7"
97 ];
98
99 disabledTestPaths = [
100 "test/extra/test_lr_scheduler.py"
101 "test/models/test_mnist.py"
102 "test/models/test_real_world.py"
103 ];
104
105 meta = with lib; {
106 description = "A simple and powerful neural network framework";
107 homepage = "https://github.com/tinygrad/tinygrad";
108 changelog = "https://github.com/tinygrad/tinygrad/releases/tag/v${version}";
109 license = licenses.mit;
110 maintainers = with maintainers; [ GaetanLepage ];
111 };
112}