Fixes

Signed-off-by: Joaquin Anton Guirao <[email protected]>

dali/test/python/operator_1/test_arithmetic_ops.py

@@ -442,8 +442,8 @@ def get_numpy_input(input, kind, orig_type, target_type):
 
 
 def extract_un_data(pipe_out, sample_id, kind, target_type):
-    input = pipe_out[0].at(sample_id).as_cpu()
-    out = pipe_out[1].at(sample_id).as_cpu()
+    input = np.array(pipe_out[0][sample_id].as_cpu())
+    out = np.array(pipe_out[1][sample_id].as_cpu())
     assert_equals(out.dtype, target_type)
     in_np = get_numpy_input(input, kind, input.dtype.type, target_type)
     return in_np, out
@@ -458,15 +458,15 @@ def extract_data(pipe_out, sample_id, kinds, target_type):
     arity = len(kinds)
     inputs = []
     for i in range(arity):
-        dali_in = pipe_out[i].at(sample_id).as_cpu()
+        dali_in = np.array(pipe_out[i][sample_id].as_cpu())
         numpy_in = get_numpy_input(
             dali_in,
             kinds[i],
             dali_in.dtype.type,
             target_type if target_type is not None else dali_in.dtype.type,
         )
         inputs.append(numpy_in)
-    out = pipe_out[arity].at(sample_id).as_cpu()
+    out = np.array(pipe_out[arity][sample_id].as_cpu())
     return tuple(inputs) + (out,)
 
 

dali/test/python/operator_1/test_coin_flip.py

@@ -50,8 +50,7 @@ def shape_gen_f():
         if shape_out is not None:
             outputs += [shape_out]
         pipe.set_outputs(*outputs)
-    outputs = pipe.run()
-    data_out = outputs[0].as_cpu()
+    (data_out,) = tuple(out.as_cpu() for out in pipe.run())
     shapes_out = None
     if max_shape is not None:
         shapes_out = outputs[1].as_cpu()

dali/test/python/operator_1/test_coord_flip.py

@@ -65,22 +65,20 @@ def check_operator_coord_flip(device, batch_size, layout, shape, center_x, cente
     eii1 = RandomDataIterator(batch_size, shape=shape, dtype=np.float32)
     pipe = CoordFlipPipeline(device, batch_size, iter(eii1), layout, center_x, center_y, center_z)
     for i in range(30):
-        outputs = pipe.run()
+        outputs0, outputs1, outputs2, outputs3, outputs4 = pipe.run()
+        outputs1 = outputs1.as_cpu()
         for sample in range(batch_size):
-            in_coords = outputs[0].at(sample)
-            if device == "gpu":
-                out_coords = outputs[1].at(sample).as_cpu()
-            else:
-                out_coords = outputs[1].at(sample)
+            in_coords = outputs0.at(sample)
+            out_coords = outputs1.at(sample)
             if in_coords.shape == () or in_coords.shape[0] == 0:
                 assert out_coords.shape == () or out_coords.shape[0] == 0
                 continue
 
-            flip_x = outputs[2].at(sample)
-            flip_y = outputs[3].at(sample)
+            flip_x = outputs2.at(sample)
+            flip_y = outputs3.at(sample)
             flip_z = None
             if len(layout) == 3:
-                flip_z = outputs[4].at(sample)
+                flip_z = outputs4.at(sample)
             _, ndim = in_coords.shape
 
             flip_dim = [flip_x, flip_y]

dali/test/python/operator_1/test_crop_mirror_normalize.py

@@ -765,11 +765,8 @@ def check_cmn_with_out_of_bounds_policy_support(
 
     if fill_values is None:
         fill_values = 0
-    for k in range(3):
-        outs = pipe.run()
-        out = outs[0]
-        in_data = outs[1]
-        mirror_data = outs[2]
+    for _ in range(3):
+        out, in_data, mirror_data = pipe.run()
         out = out.as_cpu()
         in_data = in_data.as_cpu()
 
@@ -898,13 +895,12 @@ def pipe():
     batch_size = 10
     p = pipe(batch_size=batch_size)
     for _ in range(3):
-        outs = p.run()
-        for s in range(batch_size):
-            out, image_like, mean, std = [np.array(o[s].as_cpu()) for o in outs]
+        out, image_like, mean, std = tuple(out.as_cpu() for out in p.run())
         ref_scale = scale or 1.0
         ref_shift = shift or 0.0
         ref_out = ref_scale * (image_like - mean) / std + ref_shift
-        np.testing.assert_allclose(out, ref_out, atol=ref_scale * 1e-6)
+        for s in range(batch_size):
+            np.testing.assert_allclose(out.at(s), ref_out.at(s), atol=ref_scale * 1e-6)
 
 
 def test_per_sample_norm_args():

dali/test/python/operator_1/test_input_promotion.py

@@ -41,9 +41,9 @@ def test_slice_fn():
     out_cpu = fn.slice(src, np.array([1, 1]), np.array([2, 1]), axes=[0, 1])
     out_gpu = fn.slice(src.gpu(), np.array([1, 1]), np.array([2, 1]), axes=[0, 1])
     pipe.set_outputs(out_cpu, out_gpu)
-    o = pipe.run()
-    assert np.array_equal(o[0].at(0), np.array([[14], [17]]))
-    assert np.array_equal(o[1].at(0).as_cpu(), np.array([[14], [17]]))
+    out0, out1 = tuple(out.as_cpu() for out in pipe.run())
+    assert np.array_equal(out0.at(0), np.array([[14], [17]]))
+    assert np.array_equal(np.array(out1.at(0)), np.array([[14], [17]]))
 
 
 def test_slice_ops():
@@ -56,9 +56,9 @@ def test_slice_ops():
     out_cpu = slice_cpu(src, np.array([1, 1]), np.array([2, 1]))
     out_gpu = slice_gpu(src.gpu(), np.array([1, 1]), np.array([2, 1]))
     pipe.set_outputs(out_cpu, out_gpu)
-    o = pipe.run()
-    assert np.array_equal(o[0].at(0), np.array([[14], [17]]))
-    assert np.array_equal(o[1].at(0).as_cpu(), np.array([[14], [17]]))
+    out0, out1 = tuple(out.as_cpu() for out in pipe.run())
+    assert np.array_equal(out0.at(0), np.array([[14], [17]]))
+    assert np.array_equal(out1.at(0), np.array([[14], [17]]))
 
 
 def test_python_function():

dali/test/python/operator_1/test_pad.py

@@ -98,25 +98,22 @@ def check_pad(device, batch_size, input_max_shape, axes, axis_names, align, shap
         align = [align[0] for _ in actual_axes]
     assert len(align) == len(actual_axes)
 
-    for k in range(5):
-        out1, out2 = pipe.run()
-
-        out1_data = out1.as_cpu()
+    for _ in range(5):
+        out0, out1 = tuple(out.as_cpu() for out in pipe.run())
         max_shape = [-1] * len(input_max_shape)
 
         for i in range(len(actual_axes)):
             dim = actual_axes[i]
             align_val = align[i]
             shape_arg_val = shape_arg[i]
             for i in range(batch_size):
-                input_shape = out1_data.at(i).shape
+                input_shape = out0.at(i).shape
                 if input_shape[dim] > max_shape[dim]:
                     max_shape[dim] = input_shape[dim]
 
-        out2_data = out2.as_cpu()
         for i in range(batch_size):
-            input_shape = out1_data.at(i).shape
-            output_shape = out2_data.at(i).shape
+            input_shape = out0.at(i).shape
+            output_shape = out1.at(i).shape
 
             for j in range(len(actual_axes)):
                 dim = actual_axes[j]
@@ -221,7 +218,7 @@ def check_pad_per_sample_shapes_and_alignment(device="cpu", batch_size=3, ndim=2
             in_shape, in_data, req_shape, req_align, out_pad_shape, out_pad_align, out_pad_both
         )
     for _ in range(num_iter):
-        outs = [out.as_cpu() for out in pipe.run()]
+        outs = tuple(out.as_cpu() for out in pipe.run())
         for i in range(batch_size):
             in_shape, in_data, req_shape, req_align, out_pad_shape, out_pad_align, out_pad_both = [
                 outs[out_idx].at(i) for out_idx in range(len(outs))
@@ -260,7 +257,7 @@ def check_pad_to_square(device="cpu", batch_size=3, ndim=2, num_iter=3):
         out_data = fn.pad(in_data, axis_names="HW", shape=fn.cat(side, side, axis=0))
         pipe.set_outputs(in_data, out_data)
     for _ in range(num_iter):
-        outs = [out.as_cpu() for out in pipe.run()]
+        outs = tuple(out.as_cpu() for out in pipe.run())
         for i in range(batch_size):
             in_data, out_data = [outs[out_idx].at(i) for out_idx in range(len(outs))]
             in_shape = in_data.shape

dali/test/python/operator_2/test_python_function.py

@@ -27,7 +27,7 @@
 from PIL import Image, ImageEnhance
 from nvidia.dali.ops import _DataNode
 from nose2.tools import params
-
+import numpy as np
 from nose_utils import raises
 from test_utils import get_dali_extra_path, np_type_to_dali
 
@@ -312,7 +312,9 @@ def test_python_operator_brightness():
         (numpy_output,) = numpy_brightness.run()
         (dali_output,) = dali_brightness.run()
         for i in range(len(dali_output)):
-            assert numpy.allclose(numpy_output.at(i), dali_output.at(i).as_cpu(), rtol=1e-5, atol=1)
+            assert numpy.allclose(
+                numpy_output.at(i), np.array(dali_output.at(i).as_cpu()), rtol=1e-5, atol=1
+            )
 
 
 def invalid_function(image):

dali/test/python/operator_2/test_remap.py

@@ -168,16 +168,14 @@ def _compare_pipelines_pixelwise(self, pipe1, pipe2, N_iterations, eps=0.01):
         pipe1.build()
         pipe2.build()
         for _ in range(N_iterations):
-            out1 = pipe1.run()
-            out2 = pipe2.run()
+            out1 = tuple(out.as_cpu() for out in pipe1.run())
+            out2 = tuple(out.as_cpu() for out in pipe2.run())
             self.assertTrue(
                 len(out1) == len(out2),
                 f"Numbers of outputs in the pipelines does not match: {len(out1)} vs {len(out2)}.",
             )
             for i in range(len(out1)):
-                out1_data = out1[i].as_cpu()
-                out2_data = out2[i].as_cpu()
-                for sample1, sample2 in zip(out1_data, out2_data):
+                for sample1, sample2 in zip(out1[i], out2[i]):
                     s1 = np.array(sample1)
                     s2 = np.array(sample2)
                     self.assertTrue(

dali/test/python/operator_2/test_resize.py

@@ -864,7 +864,7 @@ def resize_pipe():
 
     pipe = resize_pipe()
     (outs,) = pipe.run()
-    out = outs.at(0).as_cpu()
+    out = np.array(outs.at(0).as_cpu())
     global large_data_resized
     if large_data_resized is None:
         large_data_resized = make_cube(350, 224, 224)

dali/test/python/operator_2/test_resize_seq.py

@@ -41,7 +41,7 @@ def init_video_data():
         video_pipe.set_outputs(input)
 
     out = video_pipe.run()
-    in_seq = out[0].at(0).as_cpu()
+    in_seq = np.array(out[0].at(0).as_cpu())
     return in_seq
 
 

dali/test/python/operator_2/test_subscript.py

@@ -33,22 +33,22 @@ def test_plain_indexing():
     data = [np.float32([[0, 1, 2], [3, 4, 5]]), np.float32([[0, 1], [2, 3], [4, 5]])]
     src = fn.external_source(lambda: data, layout="AB")
     pipe = index_pipe(src, lambda x: x[1, 1])
-    inp, cpu, gpu = pipe.run()
+    inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
     for i in range(len(inp)):
         x = inp.at(i)
         assert np.array_equal(x[1, 1], cpu.at(i))
-        assert np.array_equal(x[1, 1], gpu.at(i).as_cpu())
+        assert np.array_equal(x[1, 1], np.array(gpu.at(i)))
 
 
 def _test_indexing(data_gen, input_layout, output_layout, dali_index_func, ref_index_func=None):
     src = fn.external_source(data_gen, layout=input_layout)
     pipe = index_pipe(src, dali_index_func)
-    inp, cpu, gpu = pipe.run()
+    inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
     for i in range(len(inp)):
         x = inp.at(i)
         ref = (ref_index_func or dali_index_func)(x)
         assert np.array_equal(ref, cpu.at(i))
-        assert np.array_equal(ref, gpu.at(i).as_cpu())
+        assert np.array_equal(ref, np.array(gpu.at(i)))
         assert cpu.layout() == output_layout
         assert gpu.layout() == output_layout
 
@@ -87,11 +87,11 @@ def test_swapped_ends():
     data = [np.uint8([1, 2, 3]), np.uint8([1, 2])]
     src = fn.external_source(lambda: data)
     pipe = index_pipe(src, lambda x: x[2:1])
-    inp, cpu, gpu = pipe.run()
+    inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
     for i in range(len(inp)):
         x = inp.at(i)
         assert np.array_equal(x[2:1], cpu.at(i))
-        assert np.array_equal(x[2:1], gpu.at(i).as_cpu())
+        assert np.array_equal(x[2:1], np.array(gpu.at(i)))
 
 
 def test_noop():
@@ -116,7 +116,7 @@ def data_gen():
     j = 0
     k = 0
     for _ in range(4):
-        inp, cpu, gpu = pipe.run()
+        inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
         for i in range(len(inp)):
             x = inp.at(i)
             # fmt: off
@@ -125,7 +125,7 @@ def data_gen():
             j = (j + 1) % len(lo_idxs)
             k = (k + 1) % len(hi_idxs)
             assert np.array_equal(ref, cpu.at(i))
-            assert np.array_equal(ref, gpu.at(i).as_cpu())
+            assert np.array_equal(ref, np.array(gpu.at(i)))
 
 
 def test_runtime_stride_dim1():
@@ -142,14 +142,14 @@ def data_gen():
 
     j = 0
     for _ in range(4):
-        inp, cpu, gpu = pipe.run()
+        inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
         for i in range(len(inp)):
             x = inp.at(i)
             # fmt: off
             ref = x[::strides[j]]
             # fmt: on
             assert np.array_equal(ref, cpu.at(i))
-            assert np.array_equal(ref, gpu.at(i).as_cpu())
+            assert np.array_equal(ref, np.array(gpu.at(i)))
             j = (j + 1) % len(strides)
 
 
@@ -167,14 +167,14 @@ def data_gen():
 
     j = 0
     for _ in range(4):
-        inp, cpu, gpu = pipe.run()
+        inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
         for i in range(len(inp)):
             x = inp.at(i)
             # fmt: off
             ref = x[:, ::strides[j]]
             # fmt: on
             assert np.array_equal(ref, cpu.at(i))
-            assert np.array_equal(ref, gpu.at(i).as_cpu())
+            assert np.array_equal(ref, np.array(gpu.at(i)))
             j = (j + 1) % len(strides)
 
 
@@ -287,19 +287,19 @@ def test_multiple_skipped_dims():
     ]
     src = fn.external_source(lambda: data, layout="ABCD")
     pipe = index_pipe(src, lambda x: x[1, :, :, 1])
-    inp, cpu, gpu = pipe.run()
+    inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
     for i in range(len(inp)):
         x = inp.at(i)
         assert np.array_equal(x[1, :, :, 1], cpu.at(i))
-        assert np.array_equal(x[1, :, :, 1], gpu.at(i).as_cpu())
+        assert np.array_equal(x[1, :, :, 1], np.array(gpu.at(i)))
 
 
 def test_empty_slice():
     data = [np.full((4, 5), 123), np.full((0, 1), 42)]
     src = fn.external_source(lambda: data)
     pipe = index_pipe(src, lambda x: x[0:0, 0:1])
-    inp, cpu, gpu = pipe.run()
+    inp, cpu, gpu = tuple(out.as_cpu() for out in pipe.run())
     for i in range(len(inp)):
         x = inp.at(i)
         assert np.array_equal(x[0:0, 0:1], cpu.at(i))
-        assert np.array_equal(x[0:0, 0:1], gpu.at(i).as_cpu())
+        assert np.array_equal(x[0:0, 0:1], np.array(gpu.at(i)))

dali/test/python/operator_2/test_uniform.py

@@ -22,10 +22,9 @@ def check_uniform_default(device="cpu", batch_size=32, shape=[1e5], val_range=No
     pipe = Pipeline(batch_size=batch_size, device_id=0, num_threads=3, seed=123456)
     with pipe:
         pipe.set_outputs(dali.fn.random.uniform(device=device, range=val_range, shape=shape))
-    for it in range(niter):
-        outputs = pipe.run()
+    for _ in range(niter):
+        (data_out,) = tuple(out.as_cpu() for out in pipe.run())
         val_range = (-1.0, 1.0) if val_range is None else val_range
-        data_out = outputs[0].as_cpu()
         pvs = []
         for i in range(batch_size):
             data = np.array(data_out[i])
@@ -63,9 +62,8 @@ def check_uniform_continuous_next_after(device="cpu", batch_size=32, shape=[1e5]
     pipe = Pipeline(batch_size=batch_size, device_id=0, num_threads=3, seed=123456)
     with pipe:
         pipe.set_outputs(dali.fn.random.uniform(device=device, range=val_range, shape=shape))
-    for it in range(niter):
-        outputs = pipe.run()
-        data_out = outputs[0].as_cpu()
+    for _ in range(niter):
+        (data_out,) = tuple(out.as_cpu() for out in pipe.run())
         for i in range(batch_size):
             data = np.array(data_out[i])
             assert (val_range[0] == data).all(), f"{data} is outside of requested range"
@@ -83,9 +81,8 @@ def check_uniform_discrete(device="cpu", batch_size=32, shape=[1e5], values=None
     pipe = Pipeline(batch_size=batch_size, device_id=0, num_threads=3, seed=123456)
     with pipe:
         pipe.set_outputs(dali.fn.random.uniform(device=device, values=values, shape=shape))
-    for it in range(niter):
-        outputs = pipe.run()
-        data_out = outputs[0].as_cpu()
+    for _ in range(niter):
+        (data_out,) = tuple(out.as_cpu() for out in pipe.run())
         values_set = set(values)
         maxval = np.max(values)
         bins = np.concatenate([values, np.array([np.nextafter(maxval, maxval + 1)])])