diff --git a/pkgs/development/python-modules/tinygrad/default.nix b/pkgs/development/python-modules/tinygrad/default.nix index 8566c04271122ce..bb0c6065f0a6bdc 100644 --- a/pkgs/development/python-modules/tinygrad/default.nix +++ b/pkgs/development/python-modules/tinygrad/default.nix @@ -74,6 +74,13 @@ buildPythonPackage rec { substituteInPlace tinygrad/runtime/ops_clang.py \ --replace-fail "'clang'" "'${lib.getExe clang}'" '' + # `cuda_fp16.h` and co. are needed at runtime to compile kernels + + lib.optionalString cudaSupport '' + substituteInPlace tinygrad/runtime/support/compiler_cuda.py \ + --replace-fail \ + ', "-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include/"' \ + ', "-I${lib.getDev cudaPackages.cuda_cudart}/include/"' + '' + lib.optionalString rocmSupport '' substituteInPlace tinygrad/runtime/autogen/hip.py \ --replace-fail "/opt/rocm/lib/libamdhip64.so" "${rocmPackages.clr}/lib/libamdhip64.so" \