Chaining shape-agnostic class members

Hi,

I’m doing my first steps using concrete and stumbled upon the following situation while implementing a small multi-layer perceptron. When the first and the second layer of the MLP have the same input shape everything works fine. If I try to increase the input size of the first layer, I get a value error from an operand mismatch in the matmul operation in the second layer. To me, it seems like the forward member of the Layer class is just traced with the input dimensions of the first layer. Is there a way, to enforce a tracing of the forward function with both input-shapes? Or is there another good solution to this problem?

Thank you for your help and this cool open-source library!
Best Jonas

from concrete import fhe
import numpy as np

class MLP:
    def __init__(self):
        self.layers = []

    def add(self, layer):
        self.layers.append(layer)

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        for layer in self.layers:
            inputs = layer.forward(inputs)
        return inputs

    def compile(self):
        f = lambda inputs: self.forward(inputs)
        compiler = fhe.Compiler(f, {"inputs": "encrypted"})
        inputset = [(np.random.randint(-128, 127, (1, self.layers[-1].dim[1]))) for i in range(100)]
        self.circuit = compiler.compile(inputset)
        return self.circuit

class Layer:
    def __init__(self, n_inputs: int, n_neurons: int):
        self.dim = (n_neurons, n_inputs)
        ones = [[1] * n_inputs] * n_neurons
        self.weights = np.array(ones)

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        out = np.matmul(inputs, self.weights.T)
        return out

input_size = 4

model = MLP()
model.add(Layer(input_size, 2))
model.add(Layer(2, 2))

print("CLEARTEXT")

inputs = np.array([[1] * input_size])
print("input:")
print(inputs)

output = model.forward(inputs)
print("output:")
print(output)

print("FHE")

circuit = model.compile()
circuit.keygen()
en_inputs = circuit.encrypt(inputs)
en_output = circuit.run(en_inputs)
output = circuit.decrypt(en_output)

print(f"output:")
print(output)

print(f"circuit:")
print(circuit)

Hello

  1. If there is a shape issue: I would first check it works fine without FHE. Once it works in the clear, it should work in FHE.
  2. But first, don’t you want to have a look to GitHub - zama-ai/concrete-ml: Concrete ML: Privacy Preserving ML framework using Fully Homomorphic Encryption (FHE), built on top of Concrete, with bindings to traditional ML frameworks., instead of doing everything yourself with Concrete? It will be much simpler, a lot of things were already done for the users by the ML team, here at Zama

Cheers

Hi @benoit!

Thanks for your fast reply!

  1. It works fine without FHE, but not in FHE (see example).

If input_size = 2 it works in clear and in FHE:

CLEARTEXT
input:
[[1 1]]
output:
[[4 4]]
FHE
output:
[[4 4]]
circuit:
%0 = inputs                # EncryptedTensor<int8, shape=(1, 2)>         ∈ [-127, 126]
%1 = [[1 1] [1 1]]         # ClearTensor<uint1, shape=(2, 2)>            ∈ [1, 1]
%2 = matmul(%0, %1)        # EncryptedTensor<int9, shape=(1, 2)>         ∈ [-203, 230]
%3 = [[1 1] [1 1]]         # ClearTensor<uint1, shape=(2, 2)>            ∈ [1, 1]
%4 = matmul(%2, %3)        # EncryptedTensor<int10, shape=(1, 2)>        ∈ [-406, 460]
return %4

If input_size = 4 it works in clear but not in FHE:

CLEARTEXT
input:
[[1 1 1 1]]
output:
[[8 8]]
FHE
Traceback (most recent call last):
  File "/home/jonas/silenzio/silenzio/mlp.py", line 51, in <module>
    circuit = model.compile()
  File "/home/jonas/silenzio/silenzio/mlp.py", line 20, in compile
    self.circuit = compiler.compile(inputset)
  File "/home/jonas/.pyenv/versions/concrete/lib/python3.10/site-packages/concrete/fhe/compilation/compiler.py", line 462, in compile
    self._evaluate("Compiling", inputset)
  File "/home/jonas/.pyenv/versions/concrete/lib/python3.10/site-packages/concrete/fhe/compilation/compiler.py", line 303, in _evaluate
    self._trace(first_sample)
  File "/home/jonas/.pyenv/versions/concrete/lib/python3.10/site-packages/concrete/fhe/compilation/compiler.py", line 222, in _trace
    self.graph = Tracer.trace(self.function, parameters)
  File "/home/jonas/.pyenv/versions/concrete/lib/python3.10/site-packages/concrete/fhe/tracing/tracer.py", line 85, in trace
    output_tracers: Any = function(**arguments)
  File "/home/jonas/silenzio/silenzio/mlp.py", line 17, in <lambda>
    f = lambda inputs: self.forward(inputs)
  File "/home/jonas/silenzio/silenzio/mlp.py", line 13, in forward
    inputs = layer.forward(inputs)
  File "/home/jonas/silenzio/silenzio/mlp.py", line 30, in forward
    out = np.matmul(inputs, self.weights.T)
  File "/home/jonas/.pyenv/versions/concrete/lib/python3.10/site-packages/concrete/fhe/tracing/tracer.py", line 457, in __array_ufunc__
    return Tracer._trace_numpy_operation(ufunc, *sanitized_args, **kwargs)
  File "/home/jonas/.pyenv/versions/concrete/lib/python3.10/site-packages/concrete/fhe/tracing/tracer.py", line 415, in _trace_numpy_operation
    evaluation = operation(*sample, **kwargs)
ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 4 is different from 2)
  1. Thanks for the tip, the example is very simplified and I would like to understand how to solve the problem in Concrete.

Hi @hallojs,

The issue is here:

inputset = [(np.random.randint(-128, 127, (1, self.layers[-1].dim[1]))) for i in range(100)]
                                                          ^^

You are creating an inputset of shape (1, 2), but the first layer expects an input of shape (1, 4).

If you change it to:

inputset = [(np.random.randint(-128, 127, (1, self.layers[0].dim[1]))) for i in range(100)]

it should work :slightly_smiling_face:

Please let us know if the problem is resolved for you, thanks!

1 Like

Hi @umutsahin!

Ouch, seems like I was to focused on it working in clear… Thank you, now it works!

Best Jonas

1 Like