From b3d73a60719db2fe473533b9336751ec438c91cc Mon Sep 17 00:00:00 2001
From: shanyu <santohyunolfsssigneb@gmail.com>
Date: Fri, 25 Sep 2026 18:13:57 +0800
Subject: [PATCH] Support Conv1d and Conv3d weights in Adahessian trace

get_trace averaged the Hessian diagonal over the spatial dims only for
4D kernels. Conv1d (3D) and Conv3d (5D) weights left tmp_output
unbound. Average over dims 2..ndim-1 instead; 4D is unchanged.
---
 tests/test_adahessian_conv.py | 38 +++++++++++++++++++++++++++++++++++
 torch_optimizer/adahessian.py | 11 ++++++----
 2 files changed, 45 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_adahessian_conv.py

diff --git a/tests/test_adahessian_conv.py b/tests/test_adahessian_conv.py
new file mode 100644
index 00000000..c36c76ad
--- /dev/null
+++ b/tests/test_adahessian_conv.py
@@ -0,0 +1,38 @@
+"""Adahessian must support Conv1d/Conv2d/Conv3d weights (gh issue 500).
+
+Conv1d weights are 3D and Conv3d weights are 5D. get_trace only averaged
+the Hessian diagonal over spatial dims for 4D kernels, so any other conv
+rank raised UnboundLocalError.
+"""
+import torch
+import torch_optimizer as optim
+
+
+def _step(module, x):
+    opt = optim.Adahessian(module.parameters(), lr=0.1)
+    loss = module(x).pow(2).mean()
+    loss.backward(create_graph=True)
+    opt.step()
+
+
+def test_adahessian_conv1d():
+    torch.manual_seed(500)
+    _step(torch.nn.Conv1d(2, 2, 3), torch.randn(1, 2, 8))
+
+
+def test_adahessian_conv2d():
+    torch.manual_seed(500)
+    _step(torch.nn.Conv2d(2, 2, 3), torch.randn(1, 2, 8, 8))
+
+
+def test_adahessian_conv3d():
+    torch.manual_seed(500)
+    _step(torch.nn.Conv3d(2, 2, 2), torch.randn(1, 2, 4, 4, 4))
+
+
+def test_adahessian_conv2d_trace_unchanged():
+    # The generic spatial mean must equal the old explicit dims [2, 3].
+    hv = torch.randn(2, 2, 3, 3)
+    expected = torch.mean(hv.abs(), dim=[2, 3], keepdim=True)
+    got = torch.mean(hv.abs(), dim=list(range(2, len(hv.size()))), keepdim=True)
+    assert torch.equal(got, expected)
diff --git a/torch_optimizer/adahessian.py b/torch_optimizer/adahessian.py
index 6c836481..842857a1 100644
--- a/torch_optimizer/adahessian.py
+++ b/torch_optimizer/adahessian.py
@@ -120,11 +120,14 @@ def get_trace(self, params: Params, grads: Grads) -> List[torch.Tensor]:
                 # We use that torch.abs(hv * vi) = hv.abs()
                 tmp_output = hv.abs()
 
-            elif len(param_size) == 4:  # Conv kernel
-                # Hessian diagonal block size is 9 here: torch.sum() reduces
-                # the dim 2/3.
+            else:  # Conv kernel (Conv1d/Conv2d/Conv3d weights are 3/4/5D)
+                # Hessian diagonal block sizes are k / k*k / k*k*k here:
+                # average over the spatial dims, keeping the
+                # (out_channels, in_channels) block structure.
                 # We use that torch.abs(hv * vi) = hv.abs()
-                tmp_output = torch.mean(hv.abs(), dim=[2, 3], keepdim=True)
+                tmp_output = torch.mean(
+                    hv.abs(), dim=list(range(2, len(param_size))), keepdim=True
+                )
             hutchinson_trace.append(tmp_output)
 
         return hutchinson_trace