Function _gpu_permuted_extents_torch_case

dali/test/python/test_dltensor_operator.py:534–566

(case_name, dtype, g)

def _gpu_permuted_extents_torch_case(case_name, dtype, g):
    shapes_perms = get_permute_extents_case(case_name)
    shapes, perms = tuple(zip(*shapes_perms))
    assert len(shapes) == len(perms) == len(shapes_perms)
    input_batch = get_random_torch_batch(g, shapes, dtype)
    assert len(input_batch) == len(shapes)

    # returns permuted view of the input tensors
    def permuted_tensors(batch):
        stream = current_dali_stream()
        torch_stream = torch.cuda.ExternalStream(stream)
        with torch.cuda.stream(torch_stream):
            tensors = [torch_dlpack.from_dlpack(t) for t in batch]
            assert len(tensors) == len(perms)
            tensor_views = [t.permute(perm) for t, perm in zip(tensors, perms)]
            out = [torch_dlpack.to_dlpack(t) for t in tensor_views]
            return out

    @pipeline_def(batch_size=len(input_batch), num_threads=4, device_id=0)
    def pipeline():
        data = fn.external_source(lambda: input_batch)
        data = fn.dl_tensor_python_function(
            data.gpu(), batch_processing=True, function=permuted_tensors, synchronize_stream=False
        )
        return data

    p = pipeline()
    (out,) = p.run()

    out = [numpy.array(sample) for sample in out.as_cpu()]
    ref = [numpy.array(sample).transpose(perm) for sample, perm in zip(input_batch, perms)]

    numpy.testing.assert_equal(out, ref)


def _gpu_permuted_extents_torch_suite():