From 534cc954e5fccc8b975fed7325d1bd9863ac0c6d Mon Sep 17 00:00:00 2001
From: Madoura <madouura@gmail.com>
Date: Mon, 13 Nov 2023 09:09:03 -0600
Subject: [PATCH] python3Packages.torch: Add 'tritonSupport' option

Guard this option against 'stdenv.isDarwin', which currently is unknown whether it works or not.
---
 pkgs/development/python-modules/torch/default.nix | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index 4217b87e81e7..95d2486bfbb2 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -16,9 +16,10 @@
   filelock,
   jinja2,
   networkx,
-  openai-triton,
   sympy,
   numpy, pyyaml, cffi, click, typing-extensions,
+  # ROCm build and `torch.compile` requires `openai-triton`
+  tritonSupport ? (!stdenv.isDarwin), openai-triton,
 
   # Unit tests
   hypothesis, psutil,
@@ -377,12 +378,10 @@ in buildPythonPackage rec {
     # the following are required for tensorboard support
     pillow six future tensorboard protobuf
 
-    # ROCm build and `torch.compile` requires openai-triton
-    openai-triton
-
     # torch/csrc requires `pybind11` at runtime
     pybind11
   ]
+  ++ lib.optionals tritonSupport [ openai-triton ]
   ++ lib.optionals MPISupport [ mpi ]
   ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];