From 598bd1208144dd10618acc773f63a4927e581878 Mon Sep 17 00:00:00 2001
From: blenderfreaky <blenderfreaky@mailbox.org>
Date: Sun, 21 Dec 2025 20:33:13 +0100
Subject: [PATCH 1/3] python3Packages.causal-conv1d: add rocm support

---
 .../python-modules/causal-conv1d/default.nix  | 40 +++++++++++++++----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/pkgs/development/python-modules/causal-conv1d/default.nix b/pkgs/development/python-modules/causal-conv1d/default.nix
index ec6a888cc9ea..1865164a0cf9 100644
--- a/pkgs/development/python-modules/causal-conv1d/default.nix
+++ b/pkgs/development/python-modules/causal-conv1d/default.nix
@@ -9,9 +9,14 @@
   rocmPackages,
   config,
   cudaSupport ? config.cudaSupport,
+  rocmSupport ? config.rocmSupport,
+  rocmGpuTargets ? rocmPackages.clr.localGpuTargets or rocmPackages.clr.gpuTargets,
   which,
 }:
 
+assert cudaSupport -> !rocmSupport;
+assert rocmSupport -> !cudaSupport;
+
 buildPythonPackage rec {
   pname = "causal-conv1d";
   version = "1.5.2";
@@ -30,9 +35,9 @@ buildPythonPackage rec {
     torch
   ];
 
-  nativeBuildInputs = [ which ];
+  nativeBuildInputs = [ which ] ++ lib.optionals rocmSupport [ rocmPackages.clr ];
 
-  buildInputs = (
+  buildInputs =
     lib.optionals cudaSupport (
       with cudaPackages;
       [
@@ -44,26 +49,45 @@ buildPythonPackage rec {
         libcublas
       ]
     )
-  );
+    ++ lib.optionals rocmSupport (
+      with rocmPackages;
+      [
+        rocm-core
+        rocm-device-libs
+        rocm-runtime
+        rocm-comgr
+        hipblas
+        rocblas
+        hipcub
+        rocprim
+      ]
+    );
 
   dependencies = [
     torch
   ];
 
-  # pytest tests not enabled due to nvidia GPU dependency
+  # pytest tests not enabled due to GPU dependency
   pythonImportsCheck = [ "causal_conv1d" ];
 
   env = {
     CAUSAL_CONV1D_FORCE_BUILD = "TRUE";
   }
-  // lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; };
+  // lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; }
+  // lib.optionalAttrs rocmSupport {
+    ROCM_PATH = "${rocmPackages.clr}";
+    HIP_ARCHITECTURES = builtins.concatStringsSep "," rocmGpuTargets;
+    CPLUS_INCLUDE_PATH = lib.makeSearchPath "include" [
+      rocmPackages.hipcub
+      rocmPackages.rocprim
+    ];
+  };
 
   meta = {
     description = "Causal depthwise conv1d in CUDA with a PyTorch interface";
     homepage = "https://github.com/Dao-AILab/causal-conv1d";
     license = lib.licenses.bsd3;
-    # The package requires CUDA or ROCm, the ROCm build hasn't
-    # been completed or tested, so broken if not using cuda.
-    broken = !cudaSupport;
+    # The package requires CUDA or ROCm
+    broken = !(cudaSupport || rocmSupport);
   };
 }

From 918b32b5aac60f39d6dba91322442d27cdac7e14 Mon Sep 17 00:00:00 2001
From: blenderfreaky <blenderfreaky@mailbox.org>
Date: Sun, 21 Dec 2025 21:27:40 +0100
Subject: [PATCH 2/3] python3Packages.causal-conv1d: 1.5.2 -> 1.5.4

---
 pkgs/development/python-modules/causal-conv1d/default.nix | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/causal-conv1d/default.nix b/pkgs/development/python-modules/causal-conv1d/default.nix
index 1865164a0cf9..d4dc4cba2010 100644
--- a/pkgs/development/python-modules/causal-conv1d/default.nix
+++ b/pkgs/development/python-modules/causal-conv1d/default.nix
@@ -19,14 +19,14 @@ assert rocmSupport -> !cudaSupport;
 
 buildPythonPackage rec {
   pname = "causal-conv1d";
-  version = "1.5.2";
+  version = "1.5.4";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "Dao-AILab";
     repo = "causal-conv1d";
     tag = "v${version}";
-    hash = "sha256-B2I5QiJl0p5d1BeQcMbJBAYUb10HzqFd88QMM8Rerm0=";
+    hash = "sha256-ELuvnKP2g1I2SuaWWiibXh/oDzp4n0vXkm4oeNPOdIw=";
   };
 
   build-system = [

From b57d9f4c232cce6cffffef7b8a3f32f876264627 Mon Sep 17 00:00:00 2001
From: blenderfreaky <blenderfreaky@mailbox.org>
Date: Sat, 3 Jan 2026 14:58:48 +0100
Subject: [PATCH 3/3] python3Packages.causal-conv1d: remove assert in favor of
 meta.broken

fixup
---
 .../development/python-modules/causal-conv1d/default.nix | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/causal-conv1d/default.nix b/pkgs/development/python-modules/causal-conv1d/default.nix
index d4dc4cba2010..8902d04fc0e9 100644
--- a/pkgs/development/python-modules/causal-conv1d/default.nix
+++ b/pkgs/development/python-modules/causal-conv1d/default.nix
@@ -14,9 +14,6 @@
   which,
 }:
 
-assert cudaSupport -> !rocmSupport;
-assert rocmSupport -> !cudaSupport;
-
 buildPythonPackage rec {
   pname = "causal-conv1d";
   version = "1.5.4";
@@ -87,7 +84,9 @@ buildPythonPackage rec {
     description = "Causal depthwise conv1d in CUDA with a PyTorch interface";
     homepage = "https://github.com/Dao-AILab/causal-conv1d";
     license = lib.licenses.bsd3;
-    # The package requires CUDA or ROCm
-    broken = !(cudaSupport || rocmSupport);
+
+    # The package requires either CUDA or ROCm.
+    # It doesn't work without either, nor with both.
+    broken = cudaSupport != rocmSupport;
   };
 }