From d768d6deac1e9f951cf3142017dcb7f4a0cc7384 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 30 Jan 2025 14:55:13 +0100
Subject: [PATCH] python312Packages.torchao: init at 0.11.0

---
 .../python-modules/torchao/default.nix        | 100 ++++++++++++++++++
 pkgs/top-level/python-packages.nix            |   2 +
 2 files changed, 102 insertions(+)
 create mode 100644 pkgs/development/python-modules/torchao/default.nix

diff --git a/pkgs/development/python-modules/torchao/default.nix b/pkgs/development/python-modules/torchao/default.nix
new file mode 100644
index 000000000000..66d95e175d9b
--- /dev/null
+++ b/pkgs/development/python-modules/torchao/default.nix
@@ -0,0 +1,100 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  torch,
+
+  # tests
+  bitsandbytes,
+  expecttest,
+  fire,
+  pytest-xdist,
+  pytestCheckHook,
+  parameterized,
+  tabulate,
+  transformers,
+  unittest-xml-reporting,
+}:
+
+buildPythonPackage rec {
+  pname = "ao";
+  version = "0.11.0";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "pytorch";
+    repo = "ao";
+    tag = "v${version}";
+    hash = "sha256-CNb9xaubOmIRanLq3TM4sBbszTcVK/WFpcq/sWpof44=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    torch
+  ];
+
+  env = {
+    USE_SYSTEM_LIBS = true;
+  };
+
+  # Otherwise, the tests are loading the python module from the source instead of the installed one
+  preCheck = ''
+    rm -rf torchao
+  '';
+
+  pythonImportsCheck = [
+    "torchao"
+  ];
+
+  nativeCheckInputs = [
+    bitsandbytes
+    expecttest
+    fire
+    parameterized
+    pytest-xdist
+    pytestCheckHook
+    tabulate
+    transformers
+    unittest-xml-reporting
+  ];
+
+  disabledTests =
+    [
+      # Requires internet access
+      "test_on_dummy_distilbert"
+
+      # FileNotFoundError: [Errno 2] No such file or directory: 'checkpoints/meta-llama/Llama-2-7b-chat-hf/model.pth'
+      "test_gptq_mt"
+    ]
+    ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
+      # RuntimeError: failed to initialize QNNPACK
+      "test_smooth_linear_cpu"
+
+      # torch._inductor.exc.InductorError: LoweringException: AssertionError: Expect L1_cache_size > 0 but got 0
+      "test_int8_weight_only_quant_with_freeze_0_cpu"
+      "test_int8_weight_only_quant_with_freeze_1_cpu"
+      "test_int8_weight_only_quant_with_freeze_2_cpu"
+
+      # FileNotFoundError: [Errno 2] No such file or directory: 'test.pth'
+      "test_save_load_int4woqtensors_2_cpu"
+      "test_save_load_int8woqtensors_0_cpu"
+      "test_save_load_int8woqtensors_1_cpu"
+    ];
+
+  meta = {
+    description = "PyTorch native quantization and sparsity for training and inference";
+    homepage = "https://github.com/pytorch/ao";
+    changelog = "https://github.com/pytorch/ao/releases/tag/v${version}";
+    license = lib.licenses.bsd3;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 366ee388f55b..3b3d10934b38 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -17793,6 +17793,8 @@ self: super: with self; {
 
   torchWithoutRocm = self.torch.override { rocmSupport = false; };
 
+  torchao = callPackage ../development/python-modules/torchao { };
+
   torchaudio = callPackage ../development/python-modules/torchaudio { };
 
   torchaudio-bin = callPackage ../development/python-modules/torchaudio/bin.nix { };