def test_return_ticket_np_torch(dim):
""" Test torch --> np --> torch --> np conversion"""
max_tested_ndim = 4
# Random tensor shape
tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
# Make sure complex dimension has even shape
tensor_shape[dim] = 2 * tensor_shape[dim]
complex_tensor = torch.randn(tensor_shape)
np_array = transforms.to_numpy(complex_tensor, dim=dim)
tensor_back = transforms.from_numpy(np_array, dim=dim)
np_back = transforms.to_numpy(tensor_back, dim=dim)
# Check torch --> np --> torch
assert_allclose(complex_tensor, tensor_back)
# Check np --> torch --> np
np.testing.assert_allclose(np_array, np_back)
def test_to_numpy(np_torch_tuple, dim):
""" Test torch --> np conversion (right angles)"""
from_np, from_torch = np_torch_tuple
if dim == 0:
np_array = np.array(from_np)
torch_tensor = torch.tensor(from_torch)
elif dim == 1:
np_array = np.array([from_np])
torch_tensor = torch.tensor([from_torch])
elif dim == 2:
np_array = np.array([[from_np]])
torch_tensor = torch.tensor([[from_torch]])
else:
return
np_from_torch = transforms.to_numpy(torch_tensor, dim=dim)
np.testing.assert_allclose(np_array, np_from_torch)