def div_(self, y):
"""Divide two tensors element-wise"""
# TODO: Add test coverage for this code path (next 4 lines)
if isinstance(y, float) and int(y) == y:
y = int(y)
if is_float_tensor(y) and y.frac().eq(0).all():
y = y.long()
if isinstance(y, int) or is_int_tensor(y):
# Truncate protocol for dividing by public integers:
if comm.get().get_world_size() > 2:
wraps = self.wraps()
self.share /= y
# NOTE: The multiplication here must be split into two parts
# to avoid long out-of-bounds when y <= 2 since (2 ** 63) is
# larger than the largest long integer.
self -= wraps * 4 * (int(2 ** 62) // y)
else:
self.share /= y
return self
# Otherwise multiply by reciprocal
if isinstance(y, float):
y = torch.FloatTensor([y])
assert is_float_tensor(y), "Unsupported type for div_: %s" % type(y)
return self.mul_(y.reciprocal())
def mul(self, y):
"""Perform element-wise multiplication"""
if isinstance(y, int) or is_int_tensor(y):
result = self.clone()
result.share = self.share * y
return result
return self._arithmetic_function(y, "mul")
def __init__(self,
tensor=None,
size=None,
precision=None,
src=0,
device=None):
if src == SENTINEL:
return
assert (isinstance(src, int) and src >= 0
and src < comm.get().get_world_size()), "invalid tensor source"
if device is None and hasattr(tensor, "device"):
device = tensor.device
self.encoder = FixedPointEncoder(precision_bits=precision)
if tensor is not None:
if is_int_tensor(tensor) and precision != 0:
tensor = tensor.float()
tensor = self.encoder.encode(tensor)
tensor = tensor.to(device=device)
size = tensor.size()
# Generate psuedo-random sharing of zero (PRZS) and add source's tensor
self.share = ArithmeticSharedTensor.PRZS(size, device=device).share
if self.rank == src:
assert tensor is not None, "Source must provide a data tensor"
if hasattr(tensor, "src"):
assert (
tensor.src == src
), "Source of data tensor must match source of encryption"
self.share += tensor
def _check(self, encrypted_tensor, reference, msg, tolerance=None):
if tolerance is None:
tolerance = getattr(self, "default_tolerance", 0.05)
tensor = encrypted_tensor.get_plain_text()
# Check sizes match
self.assertTrue(tensor.size() == reference.size(), msg)
self.assertTrue(is_int_tensor(reference), "reference must be a long")
test_passed = (tensor == reference).all().item() == 1
if not test_passed:
logging.info(msg)
logging.info("Result = %s;\nreference = %s" % (tensor, reference))
self.assertTrue(test_passed, msg=msg)
def test_decrypt(self):
for tensor_type in [ArithmeticSharedTensor, BinarySharedTensor]:
tensor_name = tensor_type.__name__
tensors = self.tensors
if tensor_type == ArithmeticSharedTensor:
tensors = [t for t in tensors if is_float_tensor(t)]
else:
tensors = [t for t in tensors if is_int_tensor(t)]
encrypted_tensors = [tensor_type(tensor) for tensor in tensors]
data = zip(self.sizes, tensors, encrypted_tensors)
for size, tensor, encrypted_tensor in data:
with self.benchmark(
tensor_type=tensor_name, size=size, float=self.is_float(tensor)
) as bench:
for _ in bench.iters:
tensor = encrypted_tensor.get_plain_text()
def _check(self, encrypted_tensor, reference, msg, dst=None, tolerance=None):
if tolerance is None:
tolerance = getattr(self, "default_tolerance", 0.05)
tensor = encrypted_tensor.get_plain_text(dst=dst)
if dst is not None and dst != self.rank:
self.assertIsNone(tensor)
return
# Check sizes match
self.assertTrue(tensor.size() == reference.size(), msg)
self.assertTrue(is_int_tensor(reference), "reference must be a long")
test_passed = (tensor == reference).all().item() == 1
if not test_passed:
logging.info(msg)
logging.info("Result %s" % tensor)
logging.info("Result - Reference = %s" % (tensor - reference))
self.assertTrue(test_passed, msg=msg)
def div_(self, y):
"""Divide two tensors element-wise"""
# TODO: Add test coverage for this code path (next 4 lines)
if isinstance(y, float) and int(y) == y:
y = int(y)
if is_float_tensor(y) and y.frac().eq(0).all():
y = y.long()
if isinstance(y, int) or is_int_tensor(y):
if debug_mode():
tolerance = 1.0
tensor = self.get_plain_text()
# Truncate protocol for dividing by public integers:
if comm.get().get_world_size() > 2:
wraps = self.wraps()
self.share //= y
# NOTE: The multiplication here must be split into two parts
# to avoid long out-of-bounds when y <= 2 since (2 ** 63) is
# larger than the largest long integer.
self -= wraps * 4 * (int(2 ** 62) // y)
else:
self.share //= y
if debug_mode():
if not torch.lt(
torch.abs(self.get_plain_text() * y - tensor), tolerance
).all():
raise ValueError("Final result of division is incorrect.")
return self
# Otherwise multiply by reciprocal
if isinstance(y, float):
y = torch.tensor([y], dtype=torch.float, device=self.device)
assert is_float_tensor(y), "Unsupported type for div_: %s" % type(y)
return self.mul_(y.reciprocal())
def div_(self, y):
"""Divide two tensors element-wise"""
# TODO: Add test coverage for this code path (next 4 lines)
if isinstance(y, float) and int(y) == y:
y = int(y)
if is_float_tensor(y) and y.frac().eq(0).all():
y = y.long()
if isinstance(y, int) or is_int_tensor(y):
validate = cfg.debug.validation_mode
if validate:
tolerance = 1.0
tensor = self.get_plain_text()
# Truncate protocol for dividing by public integers:
if comm.get().get_world_size() > 2:
protocol = globals()[cfg.mpc.protocol]
protocol.truncate(self, y)
else:
self.share = self.share.div_(y, rounding_mode="trunc")
# Validate
if validate:
if not torch.lt(torch.abs(self.get_plain_text() * y - tensor),
tolerance).all():
raise ValueError("Final result of division is incorrect.")
return self
# Otherwise multiply by reciprocal
if isinstance(y, float):
y = torch.tensor([y], dtype=torch.float, device=self.device)
assert is_float_tensor(y), "Unsupported type for div_: %s" % type(y)
return self.mul_(y.reciprocal())
def mul_(self, y):
"""Perform element-wise multiplication"""
if isinstance(y, int) or is_int_tensor(y):
self.share *= y
return self
return self._arithmetic_function_(y, "mul")
def __init__(
self,
tensor=None,
size=None,
broadcast_size=False,
precision=None,
src=0,
device=None,
):
"""
Creates the shared tensor from the input `tensor` provided by party `src`.
The other parties can specify a `tensor` or `size` to determine the size
of the shared tensor object to create. In this case, all parties must
specify the same (tensor) size to prevent the party's shares from varying
in size, which leads to undefined behavior.
Alternatively, the parties can set `broadcast_size` to `True` to have the
`src` party broadcast the correct size. The parties who do not know the
tensor size beforehand can provide an empty tensor as input. This is
guaranteed to produce correct behavior but requires an additional
communication round.
The parties can also set the `precision` and `device` for their share of
the tensor. If `device` is unspecified, it is set to `tensor.device`.
"""
# do nothing if source is sentinel:
if src == SENTINEL:
return
# assertions on inputs:
assert (
isinstance(src, int) and src >= 0 and src < comm.get().get_world_size()
), "specified source party does not exist"
if self.rank == src:
assert tensor is not None, "source must provide a data tensor"
if hasattr(tensor, "src"):
assert (
tensor.src == src
), "source of data tensor must match source of encryption"
if not broadcast_size:
assert (
tensor is not None or size is not None
), "must specify tensor or size, or set broadcast_size"
# if device is unspecified, try and get it from tensor:
if device is None and tensor is not None and hasattr(tensor, "device"):
device = tensor.device
# encode the input tensor:
self.encoder = FixedPointEncoder(precision_bits=precision)
if tensor is not None:
if is_int_tensor(tensor) and precision != 0:
tensor = tensor.float()
tensor = self.encoder.encode(tensor)
tensor = tensor.to(device=device)
size = tensor.size()
# if other parties do not know tensor's size, broadcast the size:
if broadcast_size:
size = comm.get().broadcast_obj(size, src)
# generate pseudo-random zero sharing (PRZS) and add source's tensor:
self.share = ArithmeticSharedTensor.PRZS(size, device=device).share
if self.rank == src:
self.share += tensor
