Function get_permute_extents_case

dali/test/python/test_dltensor_operator.py:431–531

(case_name)

def get_permute_extents_case(case_name):
    rng = random.Random(44)

    def permuted(it):
        copy = list(it)
        rng.shuffle(copy)
        return tuple(copy)

    def permuted_extents(ndim):
        extents = list(range(ndim))
        rng.shuffle(extents)
        return tuple(extents)

    # the original extents are purposely all prime numbers
    # to test handling of unaligned tensors
    prime_images = [
        (199, 181, 3),
        (1097, 61, 5),
        (373, 373, 1),
        (107, 23, 3),
        (173, 23, 5),
        (401, 167, 5),
        (181, 401, 5),
        (181, 107, 1),
        (373, 181, 5),
        (199, 401, 3),
        (1097, 1097, 1),
        (181, 61, 1),
        (401, 61, 1),
        (373, 173, 1),
        (227, 199, 3),
        (181, 1097, 1),
    ]

    if case_name == "transpose_channels_image":
        prime_images_transposed_channel = list(zip(prime_images, [(2, 0, 1)] * len(prime_images)))
        assert len(prime_images_transposed_channel) == len(prime_images)
        return prime_images_transposed_channel

    if case_name == "transpose_hw_image":
        prime_images_transposed_hw = list(zip(prime_images, [(1, 0, 2)] * len(prime_images)))
        assert len(prime_images_transposed_hw) == len(prime_images)
        return prime_images_transposed_hw

    if case_name == "image_random_permutation":
        prime_images_rnd_permuted = list(
            zip(prime_images, [permuted_extents(3) for _ in range(len(prime_images))])
        )
        assert len(prime_images_rnd_permuted) == len(prime_images)
        return prime_images_rnd_permuted

    if case_name == "transpose_channels_video":
        prime_vid_like = [
            (13, 199, 181, 3),
            (3, 1097, 61, 5),
            (17, 373, 373, 1),
            (5, 107, 23, 3),
            (11, 173, 23, 5),