neurips_tfe_matmul.py

def matmul(prot: Pond,
           x: PondMaskedTensor,
           y: PondMaskedTensor):
    a, a0, a1, alpha_on_0, alpha_on_1 = x.unwrapped
    b, b0, b1, beta_on_0, beta_on_1 = y.unwrapped

    with tf.device(prot.crypto_producer.device_name):
        ab = a.matmul(b)
        ab0, ab1 = prot._share(ab)

    with tf.device(prot.server_0.device_name):
        alpha = alpha_on_0
        beta = beta_on_0
        z0 = (ab0
              + a0.matmul(beta)
              + alpha.matmul(b0)
              + alpha.matmul(beta))

    with tf.device(prot.server_1.device_name):
        alpha = alpha_on_1
        beta = beta_on_1
        z1 = (ab1
              + a1.matmul(beta)
              + alpha.matmul(b1))

    return PondPrivateTensor(prot, z0, z1)