diff --git a/pkgs/development/python-modules/tinygrad/default.nix b/pkgs/development/python-modules/tinygrad/default.nix index e979b7991281..dabf3d10fe75 100644 --- a/pkgs/development/python-modules/tinygrad/default.nix +++ b/pkgs/development/python-modules/tinygrad/default.nix @@ -1,16 +1,17 @@ { lib, + stdenv, config, buildPythonPackage, fetchFromGitHub, + + # patches replaceVars, addDriverRunpath, - cudaSupport ? config.cudaSupport, - rocmSupport ? config.rocmSupport, cudaPackages, + llvmPackages, ocl-icd, rocmPackages, - stdenv, # build-system setuptools, @@ -21,18 +22,22 @@ unicorn, # tests + pytestCheckHook, + writableTmpDirAsHomeHook, blobfile, bottle, + capstone, clang, hexdump, hypothesis, + jax, librosa, networkx, numpy, onnx, + onnxruntime, pillow, pytest-xdist, - pytestCheckHook, safetensors, sentencepiece, tiktoken, @@ -40,19 +45,23 @@ tqdm, transformers, + # passthru tinygrad, + + cudaSupport ? config.cudaSupport, + rocmSupport ? config.rocmSupport, }: buildPythonPackage rec { pname = "tinygrad"; - version = "0.10.0"; + version = "0.10.2"; pyproject = true; src = fetchFromGitHub { owner = "tinygrad"; repo = "tinygrad"; tag = "v${version}"; - hash = "sha256-IIyTb3jDUSEP2IXK6DLsI15E5N34Utt7xv86aTHpXf8="; + hash = "sha256-BXQMacp6QjlgsVwhp2pxEZkRylZfKQhqIh92/0dPlfg="; }; patches = [ @@ -68,14 +77,30 @@ buildPythonPackage rec { postPatch = # Patch `clang` directly in the source file + # Use the unwrapped variant to enable the "native" features currently unavailable in the sandbox '' - substituteInPlace tinygrad/runtime/ops_clang.py \ - --replace-fail "'clang'" "'${lib.getExe clang}'" + substituteInPlace tinygrad/runtime/ops_cpu.py \ + --replace-fail "getenv(\"CC\", 'clang')" "'${lib.getExe llvmPackages.clang-unwrapped}'" + '' + + '' + substituteInPlace tinygrad/runtime/autogen/libc.py \ + --replace-fail "ctypes.util.find_library('c')" "'${stdenv.cc.libc}/lib/libc.so.6'" + '' + + '' + substituteInPlace tinygrad/runtime/support/llvm.py \ + --replace-fail "ctypes.util.find_library('LLVM')" "'${lib.getLib llvmPackages.llvm}/lib/libLLVM.so'" '' + lib.optionalString stdenv.hostPlatform.isLinux '' substituteInPlace tinygrad/runtime/autogen/opencl.py \ --replace-fail "ctypes.util.find_library('OpenCL')" "'${ocl-icd}/lib/libOpenCL.so'" '' + # test/test_tensor.py imports the PTX variable from the cuda_compiler.py file. + # This import leads to loading the libnvrtc.so library that is not substituted when cudaSupport = false. + # -> As a fix, we hardcode this variable to False + + lib.optionalString (!cudaSupport) '' + substituteInPlace test/test_tensor.py \ + --replace-fail "from tinygrad.runtime.support.compiler_cuda import PTX" "PTX = False" + '' # `cuda_fp16.h` and co. are needed at runtime to compile kernels + lib.optionalString cudaSupport '' substituteInPlace tinygrad/runtime/support/compiler_cuda.py \ @@ -114,18 +139,23 @@ buildPythonPackage rec { ]; nativeCheckInputs = [ + pytestCheckHook + writableTmpDirAsHomeHook + blobfile bottle + capstone clang hexdump hypothesis + jax librosa networkx numpy onnx + onnxruntime pillow pytest-xdist - pytestCheckHook safetensors sentencepiece tiktoken @@ -134,15 +164,15 @@ buildPythonPackage rec { transformers ] ++ networkx.optional-dependencies.extra; - preCheck = '' - export HOME=$(mktemp -d) - ''; - disabledTests = [ - # Fixed in https://github.com/tinygrad/tinygrad/pull/7792 - # TODO: re-enable at next release - "test_kernel_cache_in_action" + # RuntimeError: Attempting to relocate against an undefined symbol 'fmaxf' + "test_backward_sum_acc_dtype" + "test_failure_27" + + # Flaky: + # AssertionError: 2.1376906810000946 not less than 2.0 + "test_recursive_pad" # Require internet access "test_benchmark_openpilot_model" @@ -178,10 +208,6 @@ buildPythonPackage rec { "test_vgg7" ] ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-linux") [ - # Fixed in https://github.com/tinygrad/tinygrad/pull/7796 - # TODO: re-enable at next release - "test_interpolate_bilinear" - # Fail with AssertionError "test_casts_from" "test_casts_to" @@ -209,7 +235,9 @@ buildPythonPackage rec { changelog = "https://github.com/tinygrad/tinygrad/releases/tag/v${version}"; license = lib.licenses.mit; maintainers = with lib.maintainers; [ GaetanLepage ]; - # Tests segfault on darwin - badPlatforms = [ lib.systems.inspect.patterns.isDarwin ]; + badPlatforms = [ + # Tests segfault on darwin + lib.systems.inspect.patterns.isDarwin + ]; }; }