From 7a338db2eccbd60b7da3b7bed9c927117c6b3806 Mon Sep 17 00:00:00 2001
From: EtienneDosSantos <130935112+EtienneDosSantos@users.noreply.github.com>
Date: Tue, 28 May 2024 19:53:57 +0200
Subject: [PATCH] Update functional.py

---
 bitsandbytes/functional.py | 88 ++++++++++++++++++++++++++++++--------
 1 file changed, 70 insertions(+), 18 deletions(-)

diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py
index 0b1e7d5c4..bbfbf0007 100644
--- a/bitsandbytes/functional.py
+++ b/bitsandbytes/functional.py
@@ -28,57 +28,94 @@ def prod(iterable):
     """C FUNCTIONS FOR OPTIMIZERS"""
     str2optimizer32bit = {
         "adagrad": (
-            lib.cadagrad32bit_grad_32,
-            lib.cadagrad32bit_grad_16,
+            lib.cadagrad32bit_grad_fp32,
+            lib.cadagrad32bit_grad_fp16,
         ),
         "adam": (
             lib.cadam32bit_grad_fp32,
             lib.cadam32bit_grad_fp16,
             lib.cadam32bit_grad_bf16,
         ),
+        "pagedadam": (
+            lib.cpagedadam32bit_grad_fp32,
+            lib.cpagedadam32bit_grad_fp16,
+            lib.cpagedadam32bit_grad_bf16,
+        ),
+        "adamw": (
+            lib.cadam32bit_grad_fp32,
+            lib.cadam32bit_grad_fp16,
+            lib.cadam32bit_grad_bf16,
+        ),
+        "pagedadamw": (
+            lib.cpagedadam32bit_grad_fp32,
+            lib.cpagedadam32bit_grad_fp16,
+            lib.cpagedadam32bit_grad_bf16,
+        ),
         "lamb": (
             lib.cadam32bit_grad_fp32,
             lib.cadam32bit_grad_fp16,
         ),
+        "lars": (
+            lib.clars32bit_grad_fp32,
+            lib.clars32bit_grad_fp16,
+        ),
         "lion": (
             lib.clion32bit_grad_fp32,
             lib.clion32bit_grad_fp16,
             lib.clion32bit_grad_bf16,
         ),
         "momentum": (
-            lib.cmomentum32bit_grad_32,
-            lib.cmomentum32bit_grad_16,
+            lib.cmomentum32bit_grad_fp32,
+            lib.cmomentum32bit_grad_fp16,
         ),
         "rmsprop": (
-            lib.crmsprop32bit_grad_32,
-            lib.crmsprop32bit_grad_16,
+            lib.crmsprop32bit_grad_fp32,
+            lib.crmsprop32bit_grad_fp16,
         ),
     }
 
     str2optimizer8bit = {
+        "adagrad": (
+            lib.cadagrad8bit_grad_fp32,
+            lib.cadagrad8bit_grad_fp16,
+        ),
         "adam": (
-            lib.cadam_static_8bit_grad_32,
-            lib.cadam_static_8bit_grad_16,
+            lib.cadam_static_8bit_grad_fp32,
+            lib.cadam_static_8bit_grad_fp16,
+        ),
+        "pagedadam": (
+            lib.cpagedadam8bit_grad_fp32,
+            lib.cpagedadam8bit_grad_fp16,
+            lib.cpagedadam8bit_grad_bf16,
+        ),
+        "adamw": (
+            lib.cadam_static_8bit_grad_fp32,
+            lib.cadam_static_8bit_grad_fp16,
+        ),
+        "pagedadamw": (
+            lib.cpagedadam8bit_grad_fp32,
+            lib.cpagedadam8bit_grad_fp16,
+            lib.cpagedadam8bit_grad_bf16,
         ),
         "lamb": (
-            lib.cadam_static_8bit_grad_32,
-            lib.cadam_static_8bit_grad_16,
+            lib.cadam_static_8bit_grad_fp32,
+            lib.cadam_static_8bit_grad_fp16,
         ),
         "lars": (
-            lib.cmomentum_static_8bit_grad_32,
-            lib.cmomentum_static_8bit_grad_16,
+            lib.clars8bit_grad_fp32,
+            lib.clars8bit_grad_fp16,
         ),
         "lion": (
-            lib.clion_static_8bit_grad_32,
-            lib.clion_static_8bit_grad_16,
+            lib.clion_static_8bit_grad_fp32,
+            lib.clion_static_8bit_grad_fp16,
         ),
         "momentum": (
-            lib.cmomentum_static_8bit_grad_32,
-            lib.cmomentum_static_8bit_grad_16,
+            lib.cmomentum_static_8bit_grad_fp32,
+            lib.cmomentum_static_8bit_grad_fp16,
         ),
         "rmsprop": (
-            lib.crmsprop_static_8bit_grad_32,
-            lib.crmsprop_static_8bit_grad_16,
+            lib.crmsprop_static_8bit_grad_fp32,
+            lib.crmsprop_static_8bit_grad_fp16,
         ),
     }
 
@@ -92,6 +129,21 @@ def prod(iterable):
             lib.cadam_8bit_blockwise_grad_fp16,
             lib.cadam_8bit_blockwise_grad_bf16,
         ),
+        "pagedadam": (
+            lib.cpagedadam8bit_blockwise_fp32,
+            lib.cpagedadam8bit_blockwise_fp16,
+            lib.cpagedadam8bit_blockwise_bf16,
+        ),
+        "adamw": (
+            lib.cadam_8bit_blockwise_grad_fp32,
+            lib.cadam_8bit_blockwise_grad_fp16,
+            lib.cadam_8bit_blockwise_grad_bf16,
+        ),
+        "pagedadamw": (
+            lib.cpagedadam8bit_blockwise_fp32,
+            lib.cpagedadam8bit_blockwise_fp16,
+            lib.cpagedadam8bit_blockwise_bf16,
+        ),
         "lion": (
             lib.clion_8bit_blockwise_grad_fp32,
             lib.clion_8bit_blockwise_grad_fp16,