def test_dict_varying_values(self):
value1 = 3.29
value2 = 17.0
input_2D_array = np.full(shape=(4, 5),
fill_value=value1,
dtype=np.float32)
input_list = [input_2D_array, input_2D_array, input_2D_array]
input_3D_array = np.full(shape=(2, 15, 4),
fill_value=value2,
dtype=np.float64)
input_dict = dict()
input_dict["list_of_tensors"] = input_list
input_dict["single_tensor"] = input_3D_array
expected_output_2D_array = tnt.get_size() * input_2D_array
expected_output_3D_array = tnt.get_size() * input_3D_array
allreducer = tnt.TensorAllreducer(input_dict)
output_dict = allreducer.allreduce(input_dict)
assert isinstance(output_dict, dict)
assert len(output_dict) == 2
assert len(output_dict["list_of_tensors"]) == 3
assert all(
np.array_equal(array, expected_output_2D_array)
for array in output_dict["list_of_tensors"])
assert np.array_equal(output_dict["single_tensor"],
expected_output_3D_array)
def test_dict_of_tensors(self, input_dict):
expected_dict = {k: v * tnt.get_size() for k, v in input_dict.items()}
allreducer = tnt.TensorAllreducer(input_dict)
output_dict = allreducer.allreduce(input_dict)
assert {k: tf.is_tensor(v) for k, v in output_dict.items()}
assert output_dict == expected_dict
def test_single_array_identical_inputs(self, array_length, dtype):
input_array = np.ones(shape=(array_length, 1), dtype=dtype)
expected_output_array = input_array * tnt.get_size()
allreducer = tnt.TensorAllreducer(input_array)
output_array = allreducer.allreduce(input_array)
assert isinstance(output_array, np.ndarray)
assert np.array_equal(output_array, expected_output_array)
def test_tensor_from_list(self):
input_list = tf.constant([[1, 2, 3, 4, 5], [.2, .3, .4, .5, .6]])
expected_output_list = input_list * tnt.get_size()
allreducer = tnt.TensorAllreducer(input_list)
output = allreducer.allreduce(input_list)
assert tf.is_tensor(output)
assert np.all(output == expected_output_list)
def test_tensor_numeric(self, input_value, dtype):
expected_value = input_value * tnt.get_size()
input = tf.constant(input_value, dtype=dtype)
allreducer = tnt.TensorAllreducer(input)
output = allreducer.allreduce(input)
assert tf.is_tensor(output)
assert output == expected_value
def test_single_value(self):
inputs = float(tnt.get_rank())
expected_output = sum(range(tnt.get_size()))
allreducer = tnt.TensorAllreducer(inputs)
output = allreducer.allreduce(inputs)
assert isinstance(output, float)
assert expected_output == output
def test_array_inf(self, array_length, index):
injection_rank = util.same_random_int_all_ranks(0, tnt.get_size())
input_array = np.ones(shape=(array_length, 1), dtype=np.float32)
if tnt.get_rank() == injection_rank:
input_array[index] = math.inf
allreducer = tnt.TensorAllreducer(input_array)
output_array = allreducer.allreduce(input_array)
assert np.isinf(output_array[index])
def test_nd_tensor(self, input_shape, dtype):
input_array = np.ones(shape=input_shape, dtype=dtype)
expected_output_array = input_array * tnt.get_size()
inputs = tf.constant(input_array)
allreducer = tnt.TensorAllreducer(inputs)
output = allreducer.allreduce(inputs)
assert tf.is_tensor(output)
assert np.array_equal(output.numpy(), expected_output_array)
def test_dict_many_keys(self, length):
input_value = 4.2
input_dict = dict.fromkeys(("key " + str(i) for i in range(length)),
input_value)
expected_output_value = input_value * tnt.get_size()
allreducer = tnt.TensorAllreducer(input_dict)
output_dict = allreducer.allreduce(input_dict)
assert isinstance(output_dict, dict)
assert len(input_dict) == len(output_dict)
assert all(v == expected_output_value for v in output_dict.values())
def test_single_array_different_inputs(self, array_length):
input_array = np.empty(shape=(array_length, 1), dtype=np.float32)
input_array.fill(tnt.get_rank())
expected_output_array = np.empty(input_array.shape, dtype=np.float32)
expected_output_array.fill(sum(range(tnt.get_size())))
allreducer = tnt.TensorAllreducer(input_array)
output_array = allreducer.allreduce(input_array)
assert isinstance(output_array, np.ndarray)
assert np.array_equal(output_array, expected_output_array)
def test_list_of_tensors_identical_inputs(self, list_length):
input_array = tf.constant([1, 2, 3])
input_list = [input_array for i in range(list_length)]
expected_output_array = input_array * tnt.get_size()
allreducer = tnt.TensorAllreducer(input_list)
output_list = allreducer.allreduce(input_list)
assert isinstance(output_list, list)
assert all(
np.array_equal(array, expected_output_array)
for array in output_list)
def test_list_of_arrays_identical_inputs(self, list_length, dtype):
array_length = 50
input_array = np.ones(shape=(array_length, 1), dtype=dtype)
input_list = [input_array for i in range(list_length)]
expected_output_array = input_array * tnt.get_size()
allreducer = tnt.TensorAllreducer(input_list)
output_list = allreducer.allreduce(input_list)
assert isinstance(output_list, list)
assert all(
np.array_equal(array, expected_output_array)
for array in output_list)
def test_list_of_arrays_identical_inputs_diff_types(self):
input_array_float = np.ones(shape=(238, 1), dtype=np.float32)
input_array_double = np.ones(shape=(42, 1), dtype=np.double)
another_input_array_float = np.ones(shape=(99, 1), dtype=np.float32)
input_list = [
input_array_float, input_array_double, another_input_array_float
]
expected_output_list = [array * tnt.get_size() for array in input_list]
allreducer = tnt.TensorAllreducer(input_list)
output_list = allreducer.allreduce(input_list)
assert isinstance(output_list, list)
assert all(np.array_equal(output_array, expected_output_array) \
for (output_array, expected_output_array) \
in zip(output_list, expected_output_list))
def test_unsupported_type(self):
string = "sample input"
with pytest.raises(TypeError):
tnt.TensorAllreducer(string)
def create_allreducer(self, logs: Dict[str, Any]) -> None: self.allreducer = tnt.TensorAllreducer(logs, group = self.group)