diff --git a/pkgs/development/python-modules/tinygrad/default.nix b/pkgs/development/python-modules/tinygrad/default.nix index bb0c6065f0a6..ae55395b2ca2 100644 --- a/pkgs/development/python-modules/tinygrad/default.nix +++ b/pkgs/development/python-modules/tinygrad/default.nix @@ -9,14 +9,17 @@ rocmSupport ? config.rocmSupport, cudaPackages, ocl-icd, - stdenv, rocmPackages, + stdenv, # build-system setuptools, # dependencies + llvmlite, numpy, + triton, + unicorn, # tests blobfile, @@ -25,9 +28,9 @@ hexdump, hypothesis, librosa, + networkx, onnx, pillow, - pydot, pytest-xdist, pytestCheckHook, safetensors, @@ -101,6 +104,12 @@ buildPythonPackage rec { # pyobjc-framework-metal ]; + optional-dependencies = { + llvm = [ llvmlite ]; + arm = [ unicorn ]; + triton = [ triton ]; + }; + pythonImportsCheck = [ "tinygrad" @@ -116,9 +125,9 @@ buildPythonPackage rec { hexdump hypothesis librosa + networkx onnx pillow - pydot pytest-xdist pytestCheckHook safetensors @@ -127,7 +136,7 @@ buildPythonPackage rec { torch tqdm transformers - ]; + ] ++ networkx.optional-dependencies.extra; preCheck = '' export HOME=$(mktemp -d)