python3Packages.torch: fix inductor on darwin (#535046)

pkgs/development/python-modules/accelerate/default.nix

@@ -44,8 +44,6 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-IfKePiU38fUd5HefaS7J1s8Mb6hVmldINemxAJY+83o=";
   };
 
-  buildInputs = [ llvmPackages.openmp ];
-
   build-system = [ setuptools ];
 
   dependencies = [

pkgs/development/python-modules/compressed-tensors/default.nix

@@ -7,7 +7,6 @@
   # build-system
   setuptools,
   setuptools-scm,
-  llvmPackages,
 
   # dependencies
   frozendict,
@@ -47,10 +46,6 @@ buildPythonPackage (finalAttrs: {
     setuptools-scm
   ];
 
-  buildInputs = lib.optionals stdenv.cc.isClang [
-    llvmPackages.openmp
-  ];
-
   dependencies = [
     frozendict
     loguru
@@ -98,6 +93,9 @@ buildPythonPackage (finalAttrs: {
     "test_mxfp8_scales_e2e"
   ]
   ++ lib.optionals stdenv.hostPlatform.isDarwin [
+    # AssertionError: Numerical precision error
+    "test_quantization_enabled_disabled"
+
     # torch._inductor.exc.InductorError: ImportError: dlopen(/nix/var/nix/builds/nix-25002-542173852/torchinductor__nixbld1/xo/cxovsevcfanmw7lgoddbnyhoxes3nzlu7ecugxedaq2zr4f6b2qh.main.so, 0x0002):
     # symbol not found in flat namespace '___kmpc_barrier'
     "test_compress_decompress_module"

pkgs/development/python-modules/e3nn/default.nix

@@ -62,12 +62,6 @@ buildPythonPackage (finalAttrs: {
     "test_variance"
   ];
 
-  propagatedBuildInputs = lib.optionals stdenv.hostPlatform.isDarwin [
-    # Otherwise, torch will fail to include `omp.h`:
-    # torch._inductor.exc.InductorError: CppCompileError: C++ compile error OpenMP support not found
-    llvmPackages.openmp
-  ];
-
   meta = {
     description = "Modular framework for neural networks with Euclidean symmetry";
     homepage = "https://github.com/e3nn/e3nn";

pkgs/development/python-modules/hoptorch/default.nix

@@ -40,12 +40,6 @@ buildPythonPackage (finalAttrs: {
 
   pythonImportsCheck = [ "hoptorch" ];
 
-  checkInputs = lib.optionals stdenv.hostPlatform.isDarwin [
-    # torch._inductor.exc.InductorError: CppCompileError: C++ compile error
-    # fatal error: 'omp.h' file not found
-    llvmPackages.openmp
-  ];
-
   nativeCheckInputs = [
     pytestCheckHook
   ];

pkgs/development/python-modules/lerobot/default.nix

@@ -59,11 +59,6 @@ buildPythonPackage (finalAttrs: {
   ];
   dontUseCmakeConfigure = true;
 
-  checkInputs = lib.optionals stdenv.cc.isClang [
-    # test_async_iterator_* hangs after failing to find OMP headers
-    llvmPackages.openmp
-  ];
-
   pythonRelaxDeps = [
     "av"
     "diffusers"

pkgs/development/python-modules/skorch/default.nix

@@ -1,35 +1,40 @@
 {
   lib,
-  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
   numpy,
+  pandas,
   scikit-learn,
   scipy,
-  setuptools,
   tabulate,
   torch,
   tqdm,
+
+  # tests
   flaky,
-  llvmPackages,
   openssl,
-  pandas,
   pytest-cov-stub,
   pytestCheckHook,
   safetensors,
   transformers,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "skorch";
   version = "1.4.0";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "skorch-dev";
     repo = "skorch";
-    tag = "v${version}";
-    sha256 = "sha256-il3S5cfW47tKvMQGr/BfbEjMEMVzBF4gSrQhR1uKxks=";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-il3S5cfW47tKvMQGr/BfbEjMEMVzBF4gSrQhR1uKxks=";
   };
 
   build-system = [ setuptools ];
@@ -53,8 +58,6 @@ buildPythonPackage rec {
     transformers
   ];
 
-  checkInputs = lib.optionals stdenv.cc.isClang [ llvmPackages.openmp ];
-
   disabledTests = [
     # on CPU, these expect artifacts from previous GPU run
     "test_load_cuda_params_to_cpu"
@@ -77,6 +80,7 @@ buildPythonPackage rec {
     "skorch/tests/test_dataset.py"
     "skorch/tests/test_hf.py"
     "skorch/tests/llm/test_llm_classifier.py"
+
     # These tests fail when running in parallel for all platforms with:
     # "RuntimeError: The server socket has failed to listen on any local
     # network address because they use the same hardcoded port."
@@ -92,10 +96,5 @@ buildPythonPackage rec {
     changelog = "https://github.com/skorch-dev/skorch/blob/master/CHANGES.md";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [ bcdarwin ];
-    badPlatforms = [
-      # Most tests fail with:
-      # Fatal Python error: Segmentation fault
-      lib.systems.inspect.patterns.isDarwin
-    ];
   };
-}
+})

pkgs/development/python-modules/torch/source/default.nix

@@ -385,6 +385,15 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
   # until https://github.com/pytorch/pytorch/issues/76082 is addressed
   + lib.optionalString cudaSupport ''
     rm cmake/Modules/FindCUDAToolkit.cmake
+  ''
+  # Otherwise, torch compile will fail at runtime if openmp is not available
+  #   torch._inductor.exc.InductorError: CppCompileError: C++ compile error
+  #   fatal error: 'omp.h' file not found
+  + lib.optionalString stdenv.cc.isClang ''
+    substituteInPlace torch/csrc/inductor/cpp_prefix.h \
+      --replace-fail \
+        "#include <omp.h>" \
+        '#include "${lib.getInclude llvmPackages.openmp}/include/omp.h"'
   '';
 
   # NOTE(@connorbaker): Though we do not disable Gloo or MPI when building with CUDA support, caution should be taken

pkgs/development/python-modules/torchao/default.nix

@@ -77,13 +77,6 @@ buildPythonPackage (finalAttrs: {
     cpuinfo
   ];
 
-  propagatedBuildInputs = lib.optionals isDarwin [
-    # Otherwise, torch will fail to include `omp.h`:
-    # torch._inductor.exc.InductorError: CppCompileError: C++ compile error
-    # OpenMP support not found.
-    llvmPackages.openmp
-  ];
-
   dependencies = [
     torch
   ];

pkgs/development/python-modules/torchaudio/default.nix

@@ -4,9 +4,6 @@
   buildPythonPackage,
   fetchFromGitHub,
 
-  # buildInputs
-  llvmPackages,
-
   # build-system
   setuptools,
 
@@ -17,12 +14,17 @@
   # tests
   expecttest,
   inflect,
+  librosa,
   parameterized,
   pytestCheckHook,
   pytorch-lightning,
+  requests,
   scipy,
   sentencepiece,
+  soundfile,
   unidecode,
+  # linux-only:
+  fairseq,
 
   # passthru
   torchaudio,
@@ -89,13 +91,9 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
     cudaPackages.cuda_nvcc
   ];
 
-  buildInputs =
-    lib.optionals cudaSupport [
-      cudaPackages.cuda_cudart
-    ]
-    ++ lib.optionals torch.stdenv.cc.isClang [
-      llvmPackages.openmp
-    ];
+  buildInputs = lib.optionals cudaSupport [
+    cudaPackages.cuda_cudart
+  ];
 
   dependencies = [
     torch
@@ -113,11 +111,21 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
     scipy
     sentencepiece
     unidecode
+    librosa
+    requests
+    soundfile
+  ]
+  ++ lib.optionals stdenv.hostPlatform.isLinux [
+    fairseq
   ];
 
   disabledTestPaths = [
     # Require internet access
     "test/integration_tests"
+  ]
+  ++ lib.optionals stdenv.hostPlatform.isDarwin [
+    # Passes, but hangs the build after Pytest completes
+    "test/torchaudio_unittest/models/models_test.py::TestConvTasNet"
   ];
 
   disabledTests = [

pkgs/development/python-modules/torchmd-net/default.nix

@@ -22,7 +22,6 @@
   pytest-xdist,
   pytestCheckHook,
   writableTmpDirAsHomeHook,
-  llvmPackages,
 }:
 
 buildPythonPackage (finalAttrs: {
@@ -43,6 +42,13 @@ buildPythonPackage (finalAttrs: {
     setuptools-scm
   ];
 
+  env = lib.optionalAttrs stdenv.hostPlatform.isDarwin {
+    # NotImplementedError: The operator 'torchmdnet::warp_neighbor_brute_fwd' is not currently implemented for the MPS device.
+    # As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op.
+    # WARNING: this will be slower than running natively on MPS.
+    PYTORCH_ENABLE_MPS_FALLBACK = true;
+  };
+
   pythonRemoveDeps = [
     # Not a runtime dependency
     "setuptools"
@@ -66,12 +72,6 @@ buildPythonPackage (finalAttrs: {
     writableTmpDirAsHomeHook
   ];
 
-  checkInputs = lib.optionals stdenv.cc.isClang [
-    # torch._inductor.exc.InductorError: CppCompileError: C++ compile error
-    # fatal error: 'omp.h' file not found
-    llvmPackages.openmp
-  ];
-
   disabledTests = [
     # Require internet access
     "test_dataset_s66x8"