From 7a503ec0e29c96b6ef31d83c95e19670974f9791 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Mon, 7 Apr 2025 01:05:57 +0200
Subject: [PATCH] python312Packages.safetensors: add optional-dependencies

---
 .../python-modules/safetensors/default.nix    | 51 +++++++++++++++++--
 1 file changed, 48 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/safetensors/default.nix b/pkgs/development/python-modules/safetensors/default.nix
index eb7d1b502d91..e82879ccf563 100644
--- a/pkgs/development/python-modules/safetensors/default.nix
+++ b/pkgs/development/python-modules/safetensors/default.nix
@@ -5,11 +5,23 @@
   fetchFromGitHub,
   rustPlatform,
 
-  # tests
-  h5py,
+  # optional-dependencies
   numpy,
-  pytestCheckHook,
   torch,
+  tensorflow,
+  flax,
+  jax,
+  mlx,
+  paddlepaddle,
+  h5py,
+  huggingface-hub,
+  setuptools-rust,
+  pytest,
+  pytest-benchmark,
+  hypothesis,
+
+  # tests
+  pytestCheckHook,
 }:
 
 buildPythonPackage rec {
@@ -36,6 +48,39 @@ buildPythonPackage rec {
     rustPlatform.maturinBuildHook
   ];
 
+  optional-dependencies = lib.fix (self: {
+    numpy = [ numpy ];
+    torch = self.numpy ++ [
+      torch
+    ];
+    tensorflow = self.numpy ++ [
+      tensorflow
+    ];
+    pinned-tf = self.numpy ++ [
+      tensorflow
+    ];
+    jax = self.numpy ++ [
+      flax
+      jax
+    ];
+    mlx = [
+      mlx
+    ];
+    paddlepaddle = self.numpy ++ [
+      paddlepaddle
+    ];
+    testing = self.numpy ++ [
+      h5py
+      huggingface-hub
+      setuptools-rust
+      pytest
+      pytest-benchmark
+      hypothesis
+    ];
+    all = self.torch ++ self.numpy ++ self.pinned-tf ++ self.jax ++ self.paddlepaddle ++ self.testing;
+    dev = self.all;
+  });
+
   nativeCheckInputs = [
     h5py
     numpy