From d218d024ede876e328da34bf4a71d4738e495c17 Mon Sep 17 00:00:00 2001
From: David Elsing <david.elsing@posteo.net>
Date: Sat, 23 Mar 2024 22:05:06 +0000
Subject: [PATCH] gnu: Add python-pytorch-avx.

* gnu/packages/machine-learning.scm (python-pytorch-avx): New variable.
---
 gnu/packages/machine-learning.scm | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/gnu/packages/machine-learning.scm b/gnu/packages/machine-learning.scm
index d80e8076de..83fb8d4dc1 100644
--- a/gnu/packages/machine-learning.scm
+++ b/gnu/packages/machine-learning.scm
@@ -4734,6 +4734,21 @@ PyTorch when needed.
 Note: currently this package does not provide GPU support.")
     (license license:bsd-3)))
 
+;; This package variant includes the dependencies requiring at least AVX2 or
+;; AVX-512.
+(define-public python-pytorch-avx
+  (package/inherit python-pytorch
+    (name "python-pytorch-avx")
+    (inputs
+     (modify-inputs (package-inputs python-pytorch)
+       (append fbgemm nnpack)))
+    (arguments
+     (substitute-keyword-arguments (package-arguments python-pytorch)
+      ((#:phases phases)
+       #~(modify-phases #$phases
+           (delete 'disable-avx-dependencies)))))
+    (supported-systems '("x86_64-linux"))))
+
 (define-public python-pytorch-for-r-torch
   (package
     (inherit python-pytorch)