How to implement custom torch2onnx conversion ?

[MathisTLD]

I'm trying to fix auto_pad "SAME_UPPER" not being supported by implementing it myself with a wrapper.
So far here is what I do:

# this throws error operation already registered
# so we manually add the converter to the registry
# @add_converter(operation_type="Conv", version=11)
def custom_conv_converter(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
    auto_pad = node.attributes.get("auto_pad", "NOTSET")
    is_auto_pad_same = auto_pad in ("SAME_UPPER", "SAME_LOWER")
    if is_auto_pad_same:
        # remove auto_pad we will handle this ourselves
        del node._proto_attributes["auto_pad"]
        # TODO
        pass

    base_torch_module, onnx_mapping = _(node, graph)
    torch_module = base_torch_module
    if is_auto_pad_same:
        assert isinstance(base_torch_module, torch.nn.modules.conv.Conv2d)
        torch_module = OnnxConv2dWithAutoPadSame(base_torch_module, auto_pad)
        base_torch_module.padding = "valid"
        # not supported for strided conv
        # base_torch_module.padding = "same"

    return OperationConverterResult(
        torch_module=torch_module, onnx_mapping=onnx_mapping
    )


_CONVERTER_REGISTRY[OperationDescription(defs.ONNX_DOMAIN, "Conv", 11)] = (
    custom_conv_converter
)

With OnnxConv2dWithAutoPadSame being a custom wrapper

class OnnxConv2dWithAutoPadSame(nn.Module, OnnxToTorchModuleWithCustomExport):
    def __init__(
        self,
        conv: nn.Conv2d,
        mode: Literal["SAME_UPPER", "SAME_LOWER"],
    ):
        super().__init__()
        self.conv = conv
        self.mode = mode

    @override
    def _onnx_attrs(self, opset_version):
        return super()._onnx_attrs(opset_version)

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return OnnxConv2dWithAutoPadSameFunction.apply(
            input_tensor,
            self.conv.weight,
            self.conv.bias,
            self.conv.stride,
            self.conv.padding,
            self.conv.dilation,
            self.conv.groups,
            self.mode,
        )

but the breakpoint in _onnx_attrs is never reached so I try to implement it as an torch.autograd.Function

class OnnxConv2dWithAutoPadSameFunction(Function):
    @staticmethod
    def symbolic(
        ctx,
        input,
        weight,
        bias,
        stride,
        padding,
        dilation,
        groups,
        auto_pad: Literal["SAME_UPPER", "SAME_LOWER"],
    ) -> torch.Value:
        return ctx.op(
            "Conv",
            input,
            weight,
            bias,
            stride,
            padding,
            dilation,
            groups,
            auto_pad,
            outputs=1,
        )

    @staticmethod
    def forward(
        ctx,
        input,
        weight,
        bias,
        stride,
        padding,
        dilation,
        groups,
        auto_pad: Literal["SAME_UPPER", "SAME_LOWER"],
    ):
        # https://discuss.pytorch.org/t/same-padding-equivalent-in-pytorch/85121/4
        input_h, input_w = input.shape[2:]
        kh, kw = weight.shape[-2:]  # FIXME: does this also work with depth-wise conv ?
        sh, sw = stride
        # output_h = ceil(input_h / sh)
        # output_w = ceil(input_w / sh)

        if input_h % sh == 0:
            pad_y = max(kh - sh, 0)
        else:
            pad_y = max(kh - (input_h % sh), 0)
        if input_w % sw == 0:
            pad_x = max(kw - sw, 0)
        else:
            pad_x = max(kw - (input_w % sw), 0)

        pad_l = pad_x // 2  # smallest padding
        pad_r = pad_x - pad_l
        pad_t = pad_y // 2  # smallest padding
        pad_b = pad_y - pad_t

        if auto_pad == "SAME_UPPER":
            # FIXME: is this the same as spec ?
            pad_r, pad_l = pad_l, pad_r
            pad_t, pad_b = pad_b, pad_t

        padded = F.pad(
            input,
            pad=(pad_l, pad_r, pad_t, pad_b),
        )
        return torch.nn.functional.conv2d(
            padded, weight, bias, stride, padding, dilation, groups
        )

    @staticmethod
    def backward(ctx, grad_output):
        raise NotImplementedError()

Here is what I came up with but I have a ton of warnings during tracing. My symbolic method gets called but I have no idea what to return and same for the backward one (I plan on re-training)

Any help about the better way to implement my stuff ?