python3Packages.xformers: support aarch64-darwin

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

@@ -26,7 +26,8 @@
   einops,
   transformers,
   timm,
-#, flash-attn
+  #, flash-attn
+  openmp,
 }:
 let
   inherit (torch) cudaCapabilities cudaPackages cudaSupport;
@@ -66,23 +67,28 @@ buildPythonPackage {
 
   stdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv;
 
-  buildInputs = lib.optionals cudaSupport (
-    with cudaPackages;
-    [
-      # flash-attn build
-      cuda_cudart # cuda_runtime_api.h
-      libcusparse # cusparse.h
-      cuda_cccl # nv/target
-      libcublas # cublas_v2.h
-      libcusolver # cusolverDn.h
-      libcurand # curand_kernel.h
-    ]
-  );
+  buildInputs =
+    lib.optional stdenv.hostPlatform.isDarwin openmp
+    ++ lib.optionals cudaSupport (
+      with cudaPackages;
+      [
+        # flash-attn build
+        cuda_cudart # cuda_runtime_api.h
+        libcusparse # cusparse.h
+        cuda_cccl # nv/target
+        libcublas # cublas_v2.h
+        libcusolver # cusolverDn.h
+        libcurand # curand_kernel.h
+      ]
+    );
 
-  nativeBuildInputs = [
-    ninja
-    which
-  ] ++ lib.optionals cudaSupport (with cudaPackages; [ cuda_nvcc ]);
+  nativeBuildInputs =
+    [
+      ninja
+      which
+    ]
+    ++ lib.optionals cudaSupport (with cudaPackages; [ cuda_nvcc ])
+    ++ lib.optional stdenv.hostPlatform.isDarwin openmp.dev;
 
   dependencies = [
     numpy
@@ -123,9 +129,5 @@ buildPythonPackage {
     changelog = "https://github.com/facebookresearch/xformers/blob/${version}/CHANGELOG.md";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [ happysalada ];
-    badPlatforms = [
-      # fatal error: 'omp.h' file not found
-      lib.systems.inspect.patterns.isDarwin
-    ];
   };
 }

pkgs/top-level/python-packages.nix

@@ -19574,7 +19574,9 @@ self: super: with self; {
 
   xen = toPythonModule (pkgs.xen.override { python3Packages = self; });
 
-  xformers = callPackage ../development/python-modules/xformers { };
+  xformers = callPackage ../development/python-modules/xformers {
+    inherit (pkgs.llvmPackages) openmp;
+  };
 
   xgboost = callPackage ../development/python-modules/xgboost { inherit (pkgs) xgboost; };