Feature Request: support of torch.nn.Unfold function

[tguerand]

Feature request

I would like to request the support of the torch.nn.Unfold function, or a numpy equivalent.

Motivation

I am trying to do some convolutions that do different operations than regular convolutions on each kernels, therefore I would need to use a way of reorganising the images into each kernel.

I am doing it with numpy but it is quite inefficient:

def conv2D_unfold_deep(image, ratio, ksize=(1, 1), pad=0, stride=1):
    """Return the indices used in the 2D convolution"""

    nfilters, height, width = image.shape
    # arr = np.array(list(range(height * width))).reshape((image.shape))

    kernel_height, kernel_width = ksize

    output_height = int((height + 2 * pad - kernel_height) / stride) + 1
    output_width = int((width + 2 * pad - kernel_width) / stride) + 1

    out = fhe.zeros((nfilters//ratio, output_height, output_width, kernel_height * kernel_width*ratio))

    for f in range(0,nfilters,ratio):
        for h in range(output_height):
            for w in range(output_width):
                h_start = h * stride
                h_end = h_start + kernel_height
                w_start = w * stride
                w_end = w_start + kernel_width
                # Get the receptive_field
                # pad image # shape (B, 8, 10, 10)
                receptive_field = image[f:f+ratio, h_start:h_end, w_start:w_end].reshape(
                    (kernel_height * kernel_width * ratio))  # shape(B, 6) binaire
                # transform input 0/1 into int between [0 ; 2**n-1]
                out[f//ratio, h, w, :] = receptive_field  # output_var_unfold

    return out.astype(int)

which is very inefficient because of the three for loops.
I could do it with indexing but it is not supported by concrete:

import numpy as np
from concrete import fhe

coordinates = np.arange(0,100).reshape((10,10))
indexes = conv2D_unfold_deep(coordinates, ratio=1)
x = np.random.randint(0,2,(100))

@fhe.compiler({"x":"encrypted"})
def f(x):
	x = x[indexes]
	return x

f(x)

returns :

ValueError: Tracer<output=EncryptedTensor<uint1, shape=(100,)>> cannot be indexed with [[[[ 0]    [ 1]    [ 2]    [ 3]    [ 4]    [ 5]    [ 6]    [ 7]    [ 8]    [ 9]]    [[10]    [11]    [12]    [13]    [14]    [15]    [16]    [17]    [18]    [19]]    [[20]    [21]    [22]    [23]    [24]    [25]    [26]    [27]    [28]    [29]]    [[30]    [31]    [32]    [33]    [34]    [35]    [36]    [37]    [38]    [39]]    [[40]    [41]    [42]    [43]    [44]    [45]    [46]    [47]    [48]    [49]]    [[50]    [51]    [52]    [53]    [54]    [55]    [56]    [57]    [58]    [59]]    [[60]    [61]    [62]    [63]    [64]    [65]    [66]    [67]    [68]    [69]]    [[70]    [71]    [72]    [73]    [74]    [75]    [76]    [77]    [78]    [79]]    [[80]    [81]    [82]    [83]    [84]    [85]    [86]    [87]    [88]    [89]]    [[90]    [91]    [92]    [93]    [94]    [95]    [96]    [97]    [98]    [99]]]]

[umut-sahin]

Hello, here is the PR to implement fancy indexing zama-ai/concrete#640. It's ready to merge and will probably be merged soon! You'll have it in the next release of Concrete, v2.5.0 😉

[umut-sahin]

It's merged, you'll be able to do indexing in the next release 🎉