Add ConvTranspose2d with bias_jac_t_mat_prod

backpack/core/derivatives/__init__.py

@@ -1,6 +1,7 @@
 from torch.nn import (
     AvgPool2d,
     Conv2d,
+    ConvTranspose2d,
     Dropout,
     Linear,
     MaxPool2d,
@@ -12,6 +13,7 @@
 
 from .avgpool2d import AvgPool2DDerivatives
 from .conv2d import Conv2DDerivatives
+from .conv_transpose2d import ConvTranspose2DDerivatives
 from .dropout import DropoutDerivatives
 from .linear import LinearDerivatives
 from .maxpool2d import MaxPool2DDerivatives
@@ -30,4 +32,5 @@
     ReLU: ReLUDerivatives,
     Tanh: TanhDerivatives,
     Sigmoid: SigmoidDerivatives,
+    ConvTranspose2d: ConvTranspose2DDerivatives,
 }

backpack/core/derivatives/conv_transpose2d.py

@@ -0,0 +1,24 @@
+"""Partial derivatives for `torch.nn.ConvTranspose2d`."""
+
+from torch.nn import ConvTranspose2d
+
+from backpack.core.derivatives.basederivatives import BaseParameterDerivatives
+
+
+class ConvTranspose2DDerivatives(BaseParameterDerivatives):
+    def get_module(self):
+        return ConvTranspose2d
+
+    def hessian_is_zero(self):
+        return True
+
+    def _bias_jac_t_mat_prod(self, module, g_inp, g_out, mat, sum_batch=True):
+        N_axis, H_axis, W_axis = 1, 3, 4
+        axes = [H_axis, W_axis]
+        if sum_batch:
+            axes = [N_axis] + axes
+
+        return mat.sum(axes)
+
+    def _weight_jac_t_mat_prod(self, module, g_inp, g_out, mat, sum_batch=True):
+        raise NotImplementedError