def test_tensorlist_shape():
shapes = [(3, 4, 5, 6), (1, 8, 7, 6, 5), (1, ), (1, 1)]
for shape in shapes:
arr = np.empty(shape)
tl = TensorListCPU(arr)
tl_gpu = tl._as_gpu()
assert tl.shape() == [shape[1:]] * shape[0]
assert tl_gpu.shape() == [shape[1:]] * shape[0]
def test_tensorlist_dtype():
dali_types = types._all_types
np_types = list(map(dali_type_to_np, dali_types))
for dali_type, np_type in zip(dali_types, np_types):
tl = TensorListCPU([TensorCPU(np.zeros((1), dtype=np_type))])
assert tl.dtype == dali_type
assert tl._as_gpu().dtype == dali_type
def test_tl_from_list_of_tensors_same_shape():
for shape in [(10, 1), (4, 5, 6), (13, 1), (1, 1)]:
arr = np.random.rand(*shape)
tl_cpu_from_np = TensorListCPU(arr)
tl_cpu_from_tensors = TensorListCPU([TensorCPU(a) for a in arr])
np.testing.assert_array_equal(tl_cpu_from_np.as_array(),
tl_cpu_from_tensors.as_array())
tl_gpu_from_np = tl_cpu_from_np._as_gpu()
tl_gpu_from_tensors = TensorListGPU(
[TensorCPU(a)._as_gpu() for a in arr])
np.testing.assert_array_equal(tl_gpu_from_np.as_cpu().as_array(),
tl_gpu_from_tensors.as_cpu().as_array())