Feature request: support more padding method

[Soptq]

Feature request

Currently, only zero padding is supported in Concrete ML. Is it possible to also support other padding methods like reflective padding (torch.nn.ReflectionPad2d)?

Motivation

In many applications we want to make sure that the padding alters the local feature structure and global statistics the least (e.g. Image-to-image translation). Using zero padding in these applications will result in peaking edges[1][2].

Reflective padding is already supported In ONNX spec, and I think the implementation in Concrete ML should not be hard (I'm not very familiar with Concrete ML codebase, so correct me if I'm wrong). Specifically, I think we can implement a 'reflect' branch in: https://github.com/zama-ai/concrete-ml/blob/main/src/concrete/ml/onnx/onnx_impl_utils.py#L61-L66, where instead of copying the original x to the center of x_pad, we can copy different parts (9 parts: left, upper left, upper, upper right, right, bottom right, bottom, bottom left, center) of the original x to corresponding part of x_pad.

[1] https://arxiv.org/abs/1703.10593
[2] https://arxiv.org/abs/1811.11718

[andrei-stoian-zama]

It seems doable:

Currently in QuantizedPad there is:

        assert_true(
            self.mode == "constant",
            "Padding operator only supports padding with a constant",
        )

this needs to be changed to allow "reflect" from ONNX. Then in numpy_onnx_pad a new parameter for the mode should be added. The enlarged padded tensor is created here:

            x_pad = fhe_ones(tuple(padded_shape)) * numpy.int64(pad_value)`

to implement reflect one only needs to do the 9 appropriate copies from x into x_pad. doing some thing like x_pad[pads[top]:height-pads[bottom],0:pads[0]] = x[:,pads[0]:0] should be doable.