SecureNN private compare

private_compare.py

import numpy as np
import random
p = 67
bits = 64


def sample_random_tensor(shape, modulus=p):
    if len(shape) >= 1:
        n = np.prod(shape)
    else:
        n = 1
    values = [random.randrange(modulus) for _ in range(n)]
    return np.array(values).reshape(shape)


class PrivateTensor:

    def __init__(self, tensor):
        self.shares0 = sample_random_tensor(tensor.shape, p)
        self.shares1 = (tensor - self.shares0) % p


def private_compare(tensor, r, beta):

    # <variableName><number> represents which party that variable is simulating
    # for example w0, is the w vector for party 0 (j = 0)

    # i'm flipping these because the binarize function I am using is putting the most significant
    # bit at position[0], I believe the MSB should be at [63], but I have tried both ways without luck
    r_binary = np.flip(binarize(np.array([r])))
    t0 = np.flip(tensor.shares0)
    t1 = np.flip(tensor.shares1)

    w0 = np.zeros(shape=t0.shape)
    w1 = np.zeros(shape=t0.shape)
    c0 = np.zeros(shape=t1.shape)
    c1 = np.zeros(shape=t1.shape)

    for i in range(bits-1, -1, -1):

        if beta == 0:

            # party 0
            wa = (2 * r_binary[0][i] * t0[0][i])
            w0[0][i] = (t0[0][i] - wa) % p
            w0_sum = 0
            for j in range(i+1, bits):
                w0_sum += w0[0][j]
                w0_sum = w0_sum % p

            c0[0][i] = (0 - t0[0][i] + w0_sum) % p

            # party 1
            wb = (2 * r_binary[0][i] * t1[0][i])
            wbb = (r_binary[0][i] - wb)
            w1[0][i] = (t1[0][i] + wbb) % p
            w1_sum = 0
            for j in range(i+1, bits):
                w1_sum += w1[0][j]
                w1_sum = w1_sum % p

            ca = (1 + w1_sum)
            cb = (t1[0][i] + ca)
            c1[0][i] = (r_binary[0][i] - cb) % p

    answer = reconstruct(c0, c1)

    # quick check if there are any 0s in the result, if yes return 1
    for i in range(0, answer.shape[0]):
        for j in range(0, answer.shape[1]):
            if answer[0][j] == 0:
                return 1

    return 0


def reconstruct(shares0, shares1, modulus=p):
    secrets = (shares0 + shares1) % modulus
    return secrets


def binarize(tensor):
    """Computes bit decomposition of tensor

    tensor: ndarray of shape (x0, ..., xn)
    returns: a binary tensor of shape (x0, ..., xn, bits) equivalent to tensor
    """
    bitwidths = np.flip(np.arange(bits - 1, dtype=np.int64))
    for i in range(len(tensor.shape)):
        bitwidths = np.expand_dims(bitwidths, 0)
    sign = tensor < 0
    tensor = np.expand_dims(tensor, -1)
    tensor = np.right_shift(tensor, bitwidths) & 1
    tensor = np.concatenate((np.expand_dims(sign, -1), tensor), -1)
    return tensor


d = binarize(np.array([10]))
data = PrivateTensor(d)
r = 20
beta = 0

# expect result to be 10 < 20 => 0
res = private_compare(data, r, beta)
print(f'result: {res}')