def test_struct_should_forward_torch_calls_on_each_element_in_structure():
t = TensorStruct({
'points': torch.zeros((10, 2)),
'values': torch.zeros((10, 1))
})
t_expanded = t.unsqueeze(dim=1)
assert t_expanded['points'].shape == (10, 1, 2)
assert t_expanded['values'].shape == (10, 1, 1)
def test_struct_common_size_should_return_size_of_first_tensor_in_dict():
t = TensorStruct({
'a': torch.ones((10, 2)),
'b': {
'c': torch.ones((5, 2))
}
})
assert t.common_size(0) in [10, 5]
def test_struct_should_update_values_from_struct_at_given_indices():
t = TensorStruct.ones({'a': {'b': (10, 1), 'c': (10, 2)}})
new_data = TensorStruct({
'b': torch.zeros((5, 1)),
'c': torch.zeros((5, 2))
})
t['a'][2:7] = new_data
assert t['a']['b'][2, 0] == 0
assert t['a']['c'][2, 0] == 0
def test_struct_should_update_dict_if_valid_struct_given():
t = TensorStruct.ones({'a': {'b': (10, 1), 'c': (10, 2)}})
new_data = TensorStruct({
'b': torch.zeros_like(t['a']['b']),
'c': torch.zeros_like(t['a']['c'])
})
t['a'] = new_data
assert t['a']['b'][0, 0] == 0
assert t['a']['c'][0, 0] == 0
def test_struct_tensors_should_return_list_of_tensors_in_struct():
t = TensorStruct({
'a': torch.ones(5),
'b': {
'c': {
'd': torch.ones(5) * 2
}
}
})
ts = t.tensors()
assert len(ts) == 2
assert any([torch.all(torch.ones(5).eq(t_)) for t_ in ts])
assert any([torch.all(torch.ones(5).eq(t_)) * 2 for t_ in ts])
def test_struct_should_raise_if_getting_property_on_structure_with_multiple_tensors(
):
t = TensorStruct({
'points': torch.zeros((10, 2)),
'values': torch.zeros((10, 1))
})
with pytest.raises(ValueError):
_ = t.shape
def test_struct_should_return_element_when_indexing_with_string():
t = TensorStruct.ones({'a': (10, 2), 'b': {'c': (5, 4)}})
t_a = t['a']
assert isinstance(t_a, torch.Tensor)
assert t_a.shape == (10, 2)
t_b = t['b']
assert isinstance(t_b, TensorStruct)
assert t_b['c'].shape == (5, 4)
def test_struct_should_allow_to_create_nested_zeros_tensors():
t = TensorStruct.zeros(
{
'a': 5,
'b': (10, ),
'c': (3, 14),
'd': {
'e': 2,
'f': (3, 1, 4),
'g': {
'h': {
'i': (8, 2)
}
}
}
},
prefix_shape=(1, ))
td = t.data()
assert td['a'].shape == (1, 5)
assert td['b'].shape == (1, 10)
assert td['c'].shape == (1, 3, 14)
assert td['d']['e'].shape == (1, 2)
assert td['d']['f'].shape == (1, 3, 1, 4)
assert td['d']['g']['h']['i'].shape == (1, 8, 2)
def test_struct_should_raise_when_given_invalid_key():
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
with pytest.raises(KeyError):
_ = t['c']
def test_struct_should_raise_if_updating_invalid_key():
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
with pytest.raises(KeyError):
t['c'] = 5
def test_cat_should_cat_nested_tensors():
ts = [TensorStruct.ones({'a': (10, 2), 'b': (10, 3)}) for _ in range(5)]
ts_ = cat(ts, dim=0)
assert ts_['a'].shape == (50, 2)
assert ts_['b'].shape == (50, 3)
def test_struct_should_raise_if_updating_with_invalid_struct():
t = TensorStruct.ones({'a': {'b': (10, 1), 'c': (10, 2)}})
new_data = TensorStruct({'b': torch.zeros((10, 1))})
with pytest.raises(ValueError):
t['a'] = new_data
def test_struct_should_raise_if_using_unsupported_assignment():
t = TensorStruct.ones({'a': {'b': (10, 1), 'c': (10, 2)}})
with pytest.raises(ValueError):
t['a'] = 7
def test_struct_should_raise_if_constructed_from_invalid_data():
with pytest.raises(AssertionError):
_ = TensorStruct({'a': (1, 2)})
def test_struct_should_update_tensor_if_tensor_given():
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
t['a'] = torch.zeros_like(t['a'])
assert t['a'][0, 0] == 0
def test_struct_should_allow_indexing_with_list_of_indices():
t = TensorStruct.zeros({'a': (10, 2), 'b': (10, 3)})
indices = [1, 3, 7]
t_ = t[indices]
assert t_.common_size(0) == 3
def test_struct_should_return_single_elements_when_indexing_with_int():
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
t_1 = t[1]
assert isinstance(t_1, TensorStruct)
assert t_1['a'].shape == (2, )
assert t_1['b'].shape == (1, )
def test_struct_should_return_tensor_if_keep_struct_is_set_to_false():
t = TensorStruct(torch.zeros((10, 5)))
t_expanded = t.unsqueeze(dim=1, keep_struct=False)
assert isinstance(t_expanded, torch.Tensor)
assert t_expanded.shape == (10, 1, 5)
def test_struct_should_return_narrowed_elements_when_indexing_with_slice():
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
t_05 = t[:5]
assert isinstance(t_05, TensorStruct)
assert t_05['a'].shape == (5, 2)
assert t_05['b'].shape == (5, 1)
def test_struct_should_return_tensor_property_if_single_element_in_structure():
t = TensorStruct(torch.zeros((10, 5)))
assert t.shape == (10, 5)
def test_struct_should_return_narrowed_elements_when_indexing_with_tensor():
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
t_29 = t[torch.arange(2, 9, dtype=torch.long)]
assert isinstance(t_29, TensorStruct)
assert t_29['a'].shape == (7, 2)
assert t_29['b'].shape == (7, 1)
def test_struct_should_forward_torch_calls_if_single_element_in_structure():
t = TensorStruct(torch.zeros((10, 5)))
t_expanded = t.unsqueeze(dim=1)
assert t_expanded.shape == (10, 1, 5)
def test_struct_should_raise_when_given_index_other_than_string_int_slice_tensor(
):
t = TensorStruct.ones({'a': (10, 2), 'b': (10, 1)})
with pytest.raises(ValueError):
_ = t[None]
def test_tensorstruct_should_be_serializable_and_deserializable():
x = TensorStruct.zeros({
'a': (10, 2),
'b': (10, 3)
})
x_ = pickle.loads(pickle.dumps(x))
def test_struct_should_allow_to_create_single_zeros_tensor():
t = TensorStruct.zeros((2, 3), (4, 5), dtype=torch.float64, device='cpu')
assert t.shape == (4, 5, 2, 3)
assert t.dtype == torch.float64
assert t.device.type == 'cpu'