python312Packages.tinygrad: 0.10.0 -> 0.10.2 (#379553)

Author: NickCao

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,38 +22,46 @@
   unicorn,
 
   # tests
+  pytestCheckHook,
+  writableTmpDirAsHomeHook,
   blobfile,
   bottle,
+  capstone,
   clang,
   hexdump,
   hypothesis,
+  jax,
   librosa,
   networkx,
   numpy,
   onnx,
+  onnxruntime,
   pillow,
   pytest-xdist,
-  pytestCheckHook,
   safetensors,
   sentencepiece,
   tiktoken,
   torch,
   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
+    ];
   };
 }