transforms_v2_kernel_infos.py 85.1 KB
Newer Older
1
import decimal
2
3
4
5
6
import functools
import itertools
import math

import numpy as np
7
import PIL.Image
8
9
import pytest
import torch.testing
10
import torchvision.ops
11
import torchvision.transforms.v2.functional as F
12
from common_utils import (
13
    ArgsKwargs,
14
    combinations_grid,
15
16
    get_num_channels,
    ImageLoader,
17
    InfoBase,
18
    make_bounding_box_loader,
19
    make_bounding_box_loaders,
20
    make_detection_mask_loader,
21
22
    make_image_loader,
    make_image_loaders,
23
    make_image_loaders_for_interpolation,
24
    make_mask_loaders,
25
    make_video_loader,
26
    make_video_loaders,
27
28
    mark_framework_limitation,
    TestMark,
29
)
30
from torch.utils._pytree import tree_map
31
from torchvision import datapoints
32
from torchvision.transforms._functional_tensor import _max_value as get_max_value, _parse_pad_padding
33
34
35
36

__all__ = ["KernelInfo", "KERNEL_INFOS"]


37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class KernelInfo(InfoBase):
    def __init__(
        self,
        kernel,
        *,
        # Defaults to `kernel.__name__`. Should be set if the function is exposed under a different name
        # TODO: This can probably be removed after roll-out since we shouldn't have any aliasing then
        kernel_name=None,
        # Most common tests use these inputs to check the kernel. As such it should cover all valid code paths, but
        # should not include extensive parameter combinations to keep to overall test count moderate.
        sample_inputs_fn,
        # This function should mirror the kernel. It should have the same signature as the `kernel` and as such also
        # take tensors as inputs. Any conversion into another object type, e.g. PIL images or numpy arrays, should
        # happen inside the function. It should return a tensor or to be more precise an object that can be compared to
        # a tensor by `assert_close`. If omitted, no reference test will be performed.
        reference_fn=None,
        # These inputs are only used for the reference tests and thus can be comprehensive with regard to the parameter
        # values to be tested. If not specified, `sample_inputs_fn` will be used.
        reference_inputs_fn=None,
56
        # If true-ish, triggers a test that checks the kernel for consistency between uint8 and float32 inputs with the
57
        # reference inputs. This is usually used whenever we use a PIL kernel as reference.
58
59
60
61
        # Can be a callable in which case it will be called with `other_args, kwargs`. It should return the same
        # structure, but with adapted parameters. This is useful in case a parameter value is closely tied to the input
        # dtype.
        float32_vs_uint8=False,
62
63
64
        # Some kernels don't have dispatchers that would handle logging the usage. Thus, the kernel has to do it
        # manually. If set, triggers a test that makes sure this happens.
        logs_usage=False,
65
66
67
68
69
70
71
72
73
74
        # See InfoBase
        test_marks=None,
        # See InfoBase
        closeness_kwargs=None,
    ):
        super().__init__(id=kernel_name or kernel.__name__, test_marks=test_marks, closeness_kwargs=closeness_kwargs)
        self.kernel = kernel
        self.sample_inputs_fn = sample_inputs_fn
        self.reference_fn = reference_fn
        self.reference_inputs_fn = reference_inputs_fn
75

76
77
78
        if float32_vs_uint8 and not callable(float32_vs_uint8):
            float32_vs_uint8 = lambda other_args, kwargs: (other_args, kwargs)  # noqa: E731
        self.float32_vs_uint8 = float32_vs_uint8
79
        self.logs_usage = logs_usage
80
81


82
def pixel_difference_closeness_kwargs(uint8_atol, *, dtype=torch.uint8, mae=False):
83
    return dict(atol=uint8_atol / 255 * get_max_value(dtype), rtol=0, mae=mae)
84
85
86
87


def cuda_vs_cpu_pixel_difference(atol=1):
    return {
88
        (("TestKernels", "test_cuda_vs_cpu"), dtype, "cuda"): pixel_difference_closeness_kwargs(atol, dtype=dtype)
89
90
91
92
        for dtype in [torch.uint8, torch.float32]
    }


93
def pil_reference_pixel_difference(atol=1, mae=False):
94
    return {
95
        (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): pixel_difference_closeness_kwargs(
96
            atol, mae=mae
97
98
99
100
        )
    }


101
def float32_vs_uint8_pixel_difference(atol=1, mae=False):
102
103
104
105
106
    return {
        (
            ("TestKernels", "test_float32_vs_uint8"),
            torch.float32,
            "cpu",
107
        ): pixel_difference_closeness_kwargs(atol, dtype=torch.float32, mae=mae)
108
    }
109

110

111
def scripted_vs_eager_float64_tolerances(device, atol=1e-6, rtol=1e-6):
112
113
114
115
116
    return {
        (("TestKernels", "test_scripted_vs_eager"), torch.float64, device): {"atol": atol, "rtol": rtol, "mae": False},
    }


117
118
def pil_reference_wrapper(pil_kernel):
    @functools.wraps(pil_kernel)
119
120
121
122
    def wrapper(input_tensor, *other_args, **kwargs):
        if input_tensor.dtype != torch.uint8:
            raise pytest.UsageError(f"Can only test uint8 tensor images against PIL, but input is {input_tensor.dtype}")
        if input_tensor.ndim > 3:
123
            raise pytest.UsageError(
124
                f"Can only test single tensor images against PIL, but input has shape {input_tensor.shape}"
125
126
            )

127
128
129
130
131
132
133
134
135
136
137
138
139
140
        input_pil = F.to_image_pil(input_tensor)
        output_pil = pil_kernel(input_pil, *other_args, **kwargs)
        if not isinstance(output_pil, PIL.Image.Image):
            return output_pil

        output_tensor = F.to_image_tensor(output_pil)

        # 2D mask shenanigans
        if output_tensor.ndim == 2 and input_tensor.ndim == 3:
            output_tensor = output_tensor.unsqueeze(0)
        elif output_tensor.ndim == 3 and input_tensor.ndim == 2:
            output_tensor = output_tensor.squeeze(0)

        return output_tensor
141
142
143
144

    return wrapper


145
146
147
148
def xfail_jit(reason, *, condition=None):
    return TestMark(("TestKernels", "test_scripted_vs_eager"), pytest.mark.xfail(reason=reason), condition=condition)


149
def xfail_jit_python_scalar_arg(name, *, reason=None):
150
151
    return xfail_jit(
        reason or f"Python scalar int or float for `{name}` is not supported when scripting",
152
153
154
155
        condition=lambda args_kwargs: isinstance(args_kwargs.kwargs.get(name), (int, float)),
    )


156
157
158
159
KERNEL_INFOS = []


def sample_inputs_horizontal_flip_image_tensor():
160
    for image_loader in make_image_loaders(sizes=["random"], dtypes=[torch.float32]):
161
162
163
164
        yield ArgsKwargs(image_loader)


def reference_inputs_horizontal_flip_image_tensor():
165
    for image_loader in make_image_loaders(extra_dims=[()], dtypes=[torch.uint8]):
166
167
168
169
        yield ArgsKwargs(image_loader)


def sample_inputs_horizontal_flip_bounding_box():
170
    for bounding_box_loader in make_bounding_box_loaders(
171
        formats=[datapoints.BoundingBoxFormat.XYXY], dtypes=[torch.float32]
172
    ):
173
        yield ArgsKwargs(
174
            bounding_box_loader, format=bounding_box_loader.format, spatial_size=bounding_box_loader.spatial_size
175
176
177
178
        )


def sample_inputs_horizontal_flip_mask():
179
    for image_loader in make_mask_loaders(sizes=["random"], dtypes=[torch.uint8]):
180
181
182
        yield ArgsKwargs(image_loader)


183
184
185
186
187
def sample_inputs_horizontal_flip_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader)


188
189
190
191
192
193
def reference_horizontal_flip_bounding_box(bounding_box, *, format, spatial_size):
    affine_matrix = np.array(
        [
            [-1, 0, spatial_size[1]],
            [0, 1, 0],
        ],
194
        dtype="float64" if bounding_box.dtype == torch.float64 else "float32",
195
196
    )

197
198
199
    expected_bboxes = reference_affine_bounding_box_helper(
        bounding_box, format=format, spatial_size=spatial_size, affine_matrix=affine_matrix
    )
200
201
202
203
204
205
206
207
208
209
210
211
212

    return expected_bboxes


def reference_inputs_flip_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders(extra_dims=[()]):
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
            spatial_size=bounding_box_loader.spatial_size,
        )


213
214
215
216
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.horizontal_flip_image_tensor,
217
            kernel_name="horizontal_flip_image_tensor",
218
219
220
            sample_inputs_fn=sample_inputs_horizontal_flip_image_tensor,
            reference_fn=pil_reference_wrapper(F.horizontal_flip_image_pil),
            reference_inputs_fn=reference_inputs_horizontal_flip_image_tensor,
221
            float32_vs_uint8=True,
222
223
224
225
        ),
        KernelInfo(
            F.horizontal_flip_bounding_box,
            sample_inputs_fn=sample_inputs_horizontal_flip_bounding_box,
226
227
            reference_fn=reference_horizontal_flip_bounding_box,
            reference_inputs_fn=reference_inputs_flip_bounding_box,
228
229
230
231
232
        ),
        KernelInfo(
            F.horizontal_flip_mask,
            sample_inputs_fn=sample_inputs_horizontal_flip_mask,
        ),
233
234
235
236
        KernelInfo(
            F.horizontal_flip_video,
            sample_inputs_fn=sample_inputs_horizontal_flip_video,
        ),
237
238
239
240
    ]
)


241
242
243
def _get_resize_sizes(spatial_size):
    height, width = spatial_size
    length = max(spatial_size)
244
    yield length
245
246
247
248
249
    yield [length]
    yield (length,)
    new_height = int(height * 0.75)
    new_width = int(width * 1.25)
    yield [new_height, new_width]
250
251
252
    yield height, width


253
def sample_inputs_resize_image_tensor():
254
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=["RGB"], dtypes=[torch.float32]):
255
        for size in _get_resize_sizes(image_loader.spatial_size):
256
257
            yield ArgsKwargs(image_loader, size=size)

258
    for image_loader, interpolation in itertools.product(
259
        make_image_loaders(sizes=["random"], color_spaces=["RGB"]),
260
261
        [
            F.InterpolationMode.NEAREST,
262
            F.InterpolationMode.BILINEAR,
263
264
265
            F.InterpolationMode.BICUBIC,
        ],
    ):
266
        yield ArgsKwargs(image_loader, size=[min(image_loader.spatial_size) + 1], interpolation=interpolation)
267
268

    yield ArgsKwargs(make_image_loader(size=(11, 17)), size=20, max_size=25)
269
270


271
272
273
274
275
276
277
278
279
280
@pil_reference_wrapper
def reference_resize_image_tensor(*args, **kwargs):
    if not kwargs.pop("antialias", False) and kwargs.get("interpolation", F.InterpolationMode.BILINEAR) in {
        F.InterpolationMode.BILINEAR,
        F.InterpolationMode.BICUBIC,
    }:
        raise pytest.UsageError("Anti-aliasing is always active in PIL")
    return F.resize_image_pil(*args, **kwargs)


281
282
def reference_inputs_resize_image_tensor():
    for image_loader, interpolation in itertools.product(
283
        make_image_loaders_for_interpolation(),
284
285
        [
            F.InterpolationMode.NEAREST,
286
            F.InterpolationMode.NEAREST_EXACT,
287
288
289
290
            F.InterpolationMode.BILINEAR,
            F.InterpolationMode.BICUBIC,
        ],
    ):
291
        for size in _get_resize_sizes(image_loader.spatial_size):
292
293
294
295
296
297
298
299
300
301
            yield ArgsKwargs(
                image_loader,
                size=size,
                interpolation=interpolation,
                antialias=interpolation
                in {
                    F.InterpolationMode.BILINEAR,
                    F.InterpolationMode.BICUBIC,
                },
            )
302
303
304


def sample_inputs_resize_bounding_box():
305
    for bounding_box_loader in make_bounding_box_loaders():
306
        for size in _get_resize_sizes(bounding_box_loader.spatial_size):
307
            yield ArgsKwargs(bounding_box_loader, spatial_size=bounding_box_loader.spatial_size, size=size)
308
309


310
def sample_inputs_resize_mask():
311
312
    for mask_loader in make_mask_loaders(sizes=["random"], num_categories=["random"], num_objects=["random"]):
        yield ArgsKwargs(mask_loader, size=[min(mask_loader.shape[-2:]) + 1])
313
314


315
316
317
318
319
def sample_inputs_resize_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, size=[min(video_loader.shape[-2:]) + 1])


320
321
322
323
def reference_resize_bounding_box(bounding_box, *, spatial_size, size, max_size=None):
    old_height, old_width = spatial_size
    new_height, new_width = F._geometry._compute_resized_output_size(spatial_size, size=size, max_size=max_size)

324
325
326
    if (old_height, old_width) == (new_height, new_width):
        return bounding_box, (old_height, old_width)

327
328
329
330
331
    affine_matrix = np.array(
        [
            [new_width / old_width, 0, 0],
            [0, new_height / old_height, 0],
        ],
332
        dtype="float64" if bounding_box.dtype == torch.float64 else "float32",
333
334
335
    )

    expected_bboxes = reference_affine_bounding_box_helper(
336
337
338
339
        bounding_box,
        format=bounding_box.format,
        spatial_size=(new_height, new_width),
        affine_matrix=affine_matrix,
340
341
342
343
344
    )
    return expected_bboxes, (new_height, new_width)


def reference_inputs_resize_bounding_box():
345
    for bounding_box_loader in make_bounding_box_loaders(extra_dims=((), (4,))):
346
347
348
349
        for size in _get_resize_sizes(bounding_box_loader.spatial_size):
            yield ArgsKwargs(bounding_box_loader, size=size, spatial_size=bounding_box_loader.spatial_size)


350
351
352
353
354
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.resize_image_tensor,
            sample_inputs_fn=sample_inputs_resize_image_tensor,
355
            reference_fn=reference_resize_image_tensor,
356
            reference_inputs_fn=reference_inputs_resize_image_tensor,
357
            float32_vs_uint8=True,
358
            closeness_kwargs={
359
                **pil_reference_pixel_difference(10, mae=True),
360
                **cuda_vs_cpu_pixel_difference(),
361
                **float32_vs_uint8_pixel_difference(1, mae=True),
362
            },
363
            test_marks=[
364
                xfail_jit_python_scalar_arg("size"),
365
            ],
366
367
368
369
        ),
        KernelInfo(
            F.resize_bounding_box,
            sample_inputs_fn=sample_inputs_resize_bounding_box,
370
371
            reference_fn=reference_resize_bounding_box,
            reference_inputs_fn=reference_inputs_resize_bounding_box,
372
373
374
            closeness_kwargs={
                (("TestKernels", "test_against_reference"), torch.int64, "cpu"): dict(atol=1, rtol=0),
            },
375
            test_marks=[
376
                xfail_jit_python_scalar_arg("size"),
377
            ],
378
        ),
379
380
381
        KernelInfo(
            F.resize_mask,
            sample_inputs_fn=sample_inputs_resize_mask,
382
            closeness_kwargs=pil_reference_pixel_difference(10),
383
            test_marks=[
384
                xfail_jit_python_scalar_arg("size"),
385
            ],
386
        ),
387
388
389
        KernelInfo(
            F.resize_video,
            sample_inputs_fn=sample_inputs_resize_video,
390
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
391
        ),
392
393
394
395
396
397
398
399
400
401
402
403
    ]
)


_AFFINE_KWARGS = combinations_grid(
    angle=[-87, 15, 90],
    translate=[(5, 5), (-5, -5)],
    scale=[0.77, 1.27],
    shear=[(12, 12), (0, 0)],
)


404
405
406
407
408
409
410
411
412
413
def _diversify_affine_kwargs_types(affine_kwargs):
    angle = affine_kwargs["angle"]
    for diverse_angle in [int(angle), float(angle)]:
        yield dict(affine_kwargs, angle=diverse_angle)

    shear = affine_kwargs["shear"]
    for diverse_shear in [tuple(shear), list(shear), int(shear[0]), float(shear[0])]:
        yield dict(affine_kwargs, shear=diverse_shear)


414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
def _full_affine_params(**partial_params):
    partial_params.setdefault("angle", 0.0)
    partial_params.setdefault("translate", [0.0, 0.0])
    partial_params.setdefault("scale", 1.0)
    partial_params.setdefault("shear", [0.0, 0.0])
    partial_params.setdefault("center", None)
    return partial_params


_DIVERSE_AFFINE_PARAMS = [
    _full_affine_params(**{name: arg})
    for name, args in [
        ("angle", [1.0, 2]),
        ("translate", [[1.0, 0.5], [1, 2], (1.0, 0.5), (1, 2)]),
        ("scale", [0.5]),
        ("shear", [1.0, 2, [1.0], [2], (1.0,), (2,), [1.0, 0.5], [1, 2], (1.0, 0.5), (1, 2)]),
        ("center", [None, [1.0, 0.5], [1, 2], (1.0, 0.5), (1, 2)]),
    ]
    for arg in args
]


436
def get_fills(*, num_channels, dtype):
437
438
    yield None

439
440
441
442
    int_value = get_max_value(dtype)
    float_value = int_value / 2
    yield int_value
    yield float_value
443

444
445
446
    for vector_type in [list, tuple]:
        yield vector_type([int_value])
        yield vector_type([float_value])
447

448
449
450
        if num_channels > 1:
            yield vector_type(float_value * c / 10 for c in range(num_channels))
            yield vector_type(int_value if c % 2 == 0 else 0 for c in range(num_channels))
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465


def float32_vs_uint8_fill_adapter(other_args, kwargs):
    fill = kwargs.get("fill")
    if fill is None:
        return other_args, kwargs

    if isinstance(fill, (int, float)):
        fill /= 255
    else:
        fill = type(fill)(fill_ / 255 for fill_ in fill)

    return other_args, dict(kwargs, fill=fill)


466
def sample_inputs_affine_image_tensor():
467
    make_affine_image_loaders = functools.partial(
468
        make_image_loaders, sizes=["random"], color_spaces=["RGB"], dtypes=[torch.float32]
469
470
471
472
473
474
    )

    for image_loader, affine_params in itertools.product(make_affine_image_loaders(), _DIVERSE_AFFINE_PARAMS):
        yield ArgsKwargs(image_loader, **affine_params)

    for image_loader in make_affine_image_loaders():
475
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
476
477
478
479
            yield ArgsKwargs(image_loader, **_full_affine_params(), fill=fill)

    for image_loader, interpolation in itertools.product(
        make_affine_image_loaders(),
480
481
482
483
        [
            F.InterpolationMode.NEAREST,
            F.InterpolationMode.BILINEAR,
        ],
484
    ):
485
        yield ArgsKwargs(image_loader, **_full_affine_params(), fill=0)
486

487
488

def reference_inputs_affine_image_tensor():
489
    for image_loader, affine_kwargs in itertools.product(make_image_loaders_for_interpolation(), _AFFINE_KWARGS):
490
        yield ArgsKwargs(
491
            image_loader,
492
493
494
495
496
497
            interpolation=F.InterpolationMode.NEAREST,
            **affine_kwargs,
        )


def sample_inputs_affine_bounding_box():
498
    for bounding_box_loader, affine_params in itertools.product(
499
        make_bounding_box_loaders(formats=[datapoints.BoundingBoxFormat.XYXY]), _DIVERSE_AFFINE_PARAMS
500
501
502
503
    ):
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
504
            spatial_size=bounding_box_loader.spatial_size,
505
            **affine_params,
506
507
        )

508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531

def _compute_affine_matrix(angle, translate, scale, shear, center):
    rot = math.radians(angle)
    cx, cy = center
    tx, ty = translate
    sx, sy = [math.radians(sh_) for sh_ in shear]

    c_matrix = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]])
    t_matrix = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]])
    c_matrix_inv = np.linalg.inv(c_matrix)
    rs_matrix = np.array(
        [
            [scale * math.cos(rot), -scale * math.sin(rot), 0],
            [scale * math.sin(rot), scale * math.cos(rot), 0],
            [0, 0, 1],
        ]
    )
    shear_x_matrix = np.array([[1, -math.tan(sx), 0], [0, 1, 0], [0, 0, 1]])
    shear_y_matrix = np.array([[1, 0, 0], [-math.tan(sy), 1, 0], [0, 0, 1]])
    rss_matrix = np.matmul(rs_matrix, np.matmul(shear_y_matrix, shear_x_matrix))
    true_matrix = np.matmul(t_matrix, np.matmul(c_matrix, np.matmul(rss_matrix, c_matrix_inv)))
    return true_matrix


532
533
def reference_affine_bounding_box_helper(bounding_box, *, format, spatial_size, affine_matrix):
    def transform(bbox, affine_matrix_, format_, spatial_size_):
534
535
        # Go to float before converting to prevent precision loss in case of CXCYWH -> XYXY and W or H is 1
        in_dtype = bbox.dtype
536
537
        if not torch.is_floating_point(bbox):
            bbox = bbox.float()
538
        bbox_xyxy = F.convert_format_bounding_box(
539
540
541
542
            bbox.as_subclass(torch.Tensor),
            old_format=format_,
            new_format=datapoints.BoundingBoxFormat.XYXY,
            inplace=True,
543
        )
544
545
546
547
548
549
550
551
        points = np.array(
            [
                [bbox_xyxy[0].item(), bbox_xyxy[1].item(), 1.0],
                [bbox_xyxy[2].item(), bbox_xyxy[1].item(), 1.0],
                [bbox_xyxy[0].item(), bbox_xyxy[3].item(), 1.0],
                [bbox_xyxy[2].item(), bbox_xyxy[3].item(), 1.0],
            ]
        )
552
        transformed_points = np.matmul(points, affine_matrix_.T)
553
554
        out_bbox = torch.tensor(
            [
555
556
557
558
                np.min(transformed_points[:, 0]).item(),
                np.min(transformed_points[:, 1]).item(),
                np.max(transformed_points[:, 0]).item(),
                np.max(transformed_points[:, 1]).item(),
559
            ],
560
            dtype=bbox_xyxy.dtype,
561
        )
562
        out_bbox = F.convert_format_bounding_box(
563
            out_bbox, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format_, inplace=True
564
        )
565
566
567
568
        # It is important to clamp before casting, especially for CXCYWH format, dtype=int64
        out_bbox = F.clamp_bounding_box(out_bbox, format=format_, spatial_size=spatial_size_)
        out_bbox = out_bbox.to(dtype=in_dtype)
        return out_bbox
569
570
571
572

    if bounding_box.ndim < 2:
        bounding_box = [bounding_box]

573
    expected_bboxes = [transform(bbox, affine_matrix, format, spatial_size) for bbox in bounding_box]
574
575
576
577
578
579
580
581
    if len(expected_bboxes) > 1:
        expected_bboxes = torch.stack(expected_bboxes)
    else:
        expected_bboxes = expected_bboxes[0]

    return expected_bboxes


582
583
584
585
586
587
588
def reference_affine_bounding_box(bounding_box, *, format, spatial_size, angle, translate, scale, shear, center=None):
    if center is None:
        center = [s * 0.5 for s in spatial_size[::-1]]

    affine_matrix = _compute_affine_matrix(angle, translate, scale, shear, center)
    affine_matrix = affine_matrix[:2, :]

589
590
591
    expected_bboxes = reference_affine_bounding_box_helper(
        bounding_box, format=format, spatial_size=spatial_size, affine_matrix=affine_matrix
    )
592
593
594
595

    return expected_bboxes


596
def reference_inputs_affine_bounding_box():
597
598
599
    for bounding_box_loader, affine_kwargs in itertools.product(
        make_bounding_box_loaders(extra_dims=[()]),
        _AFFINE_KWARGS,
600
601
602
603
    ):
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
604
            spatial_size=bounding_box_loader.spatial_size,
605
            **affine_kwargs,
606
607
608
        )


609
def sample_inputs_affine_mask():
610
611
    for mask_loader in make_mask_loaders(sizes=["random"], num_categories=["random"], num_objects=["random"]):
        yield ArgsKwargs(mask_loader, **_full_affine_params())
612

613

614
615
616
617
618
def sample_inputs_affine_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, **_full_affine_params())


619
620
621
622
623
624
625
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.affine_image_tensor,
            sample_inputs_fn=sample_inputs_affine_image_tensor,
            reference_fn=pil_reference_wrapper(F.affine_image_pil),
            reference_inputs_fn=reference_inputs_affine_image_tensor,
626
            float32_vs_uint8=True,
627
            closeness_kwargs=pil_reference_pixel_difference(10, mae=True),
628
629
            test_marks=[
                xfail_jit_python_scalar_arg("shear"),
630
                xfail_jit_python_scalar_arg("fill"),
631
            ],
632
633
634
635
636
637
        ),
        KernelInfo(
            F.affine_bounding_box,
            sample_inputs_fn=sample_inputs_affine_bounding_box,
            reference_fn=reference_affine_bounding_box,
            reference_inputs_fn=reference_inputs_affine_bounding_box,
638
            test_marks=[
639
                xfail_jit_python_scalar_arg("shear"),
640
            ],
641
        ),
642
643
        KernelInfo(
            F.affine_mask,
644
            sample_inputs_fn=sample_inputs_affine_mask,
645
646
647
            test_marks=[
                xfail_jit_python_scalar_arg("shear"),
            ],
648
        ),
649
650
651
652
        KernelInfo(
            F.affine_video,
            sample_inputs_fn=sample_inputs_affine_video,
        ),
653
654
    ]
)
655
656
657


def sample_inputs_convert_format_bounding_box():
658
    formats = list(datapoints.BoundingBoxFormat)
659
    for bounding_box_loader, new_format in itertools.product(make_bounding_box_loaders(formats=formats), formats):
660
        yield ArgsKwargs(bounding_box_loader, old_format=bounding_box_loader.format, new_format=new_format)
661
662


663
def reference_convert_format_bounding_box(bounding_box, old_format, new_format):
664
    return torchvision.ops.box_convert(
665
666
        bounding_box, in_fmt=old_format.name.lower(), out_fmt=new_format.name.lower()
    ).to(bounding_box.dtype)
667
668
669


def reference_inputs_convert_format_bounding_box():
670
    for args_kwargs in sample_inputs_convert_format_bounding_box():
671
672
        if len(args_kwargs.args[0].shape) == 2:
            yield args_kwargs
673
674
675
676
677
678
679
680


KERNEL_INFOS.append(
    KernelInfo(
        F.convert_format_bounding_box,
        sample_inputs_fn=sample_inputs_convert_format_bounding_box,
        reference_fn=reference_convert_format_bounding_box,
        reference_inputs_fn=reference_inputs_convert_format_bounding_box,
681
        logs_usage=True,
682
683
684
685
    ),
)


686
687
688
689
690
691
def sample_inputs_vertical_flip_image_tensor():
    for image_loader in make_image_loaders(sizes=["random"], dtypes=[torch.float32]):
        yield ArgsKwargs(image_loader)


def reference_inputs_vertical_flip_image_tensor():
692
    for image_loader in make_image_loaders(extra_dims=[()], dtypes=[torch.uint8]):
693
694
695
696
697
        yield ArgsKwargs(image_loader)


def sample_inputs_vertical_flip_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders(
698
        formats=[datapoints.BoundingBoxFormat.XYXY], dtypes=[torch.float32]
699
700
    ):
        yield ArgsKwargs(
701
            bounding_box_loader, format=bounding_box_loader.format, spatial_size=bounding_box_loader.spatial_size
702
703
704
705
706
707
708
709
        )


def sample_inputs_vertical_flip_mask():
    for image_loader in make_mask_loaders(sizes=["random"], dtypes=[torch.uint8]):
        yield ArgsKwargs(image_loader)


710
711
712
713
714
def sample_inputs_vertical_flip_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader)


715
716
717
718
719
720
def reference_vertical_flip_bounding_box(bounding_box, *, format, spatial_size):
    affine_matrix = np.array(
        [
            [1, 0, 0],
            [0, -1, spatial_size[0]],
        ],
721
        dtype="float64" if bounding_box.dtype == torch.float64 else "float32",
722
723
    )

724
725
726
    expected_bboxes = reference_affine_bounding_box_helper(
        bounding_box, format=format, spatial_size=spatial_size, affine_matrix=affine_matrix
    )
727
728
729
730

    return expected_bboxes


731
732
733
734
735
736
737
738
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.vertical_flip_image_tensor,
            kernel_name="vertical_flip_image_tensor",
            sample_inputs_fn=sample_inputs_vertical_flip_image_tensor,
            reference_fn=pil_reference_wrapper(F.vertical_flip_image_pil),
            reference_inputs_fn=reference_inputs_vertical_flip_image_tensor,
739
            float32_vs_uint8=True,
740
741
742
743
        ),
        KernelInfo(
            F.vertical_flip_bounding_box,
            sample_inputs_fn=sample_inputs_vertical_flip_bounding_box,
744
745
            reference_fn=reference_vertical_flip_bounding_box,
            reference_inputs_fn=reference_inputs_flip_bounding_box,
746
747
748
749
750
        ),
        KernelInfo(
            F.vertical_flip_mask,
            sample_inputs_fn=sample_inputs_vertical_flip_mask,
        ),
751
752
753
754
        KernelInfo(
            F.vertical_flip_video,
            sample_inputs_fn=sample_inputs_vertical_flip_video,
        ),
755
756
757
758
759
760
761
    ]
)

_ROTATE_ANGLES = [-87, 15, 90]


def sample_inputs_rotate_image_tensor():
762
    make_rotate_image_loaders = functools.partial(
763
        make_image_loaders, sizes=["random"], color_spaces=["RGB"], dtypes=[torch.float32]
764
765
766
767
768
769
770
    )

    for image_loader in make_rotate_image_loaders():
        yield ArgsKwargs(image_loader, angle=15.0, expand=True)

    for image_loader, center in itertools.product(
        make_rotate_image_loaders(), [None, [1.0, 0.5], [1, 2], (1.0, 0.5), (1, 2)]
771
    ):
772
        yield ArgsKwargs(image_loader, angle=15.0, center=center)
773

774
    for image_loader in make_rotate_image_loaders():
775
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
776
777
778
779
780
781
782
            yield ArgsKwargs(image_loader, angle=15.0, fill=fill)

    for image_loader, interpolation in itertools.product(
        make_rotate_image_loaders(),
        [F.InterpolationMode.NEAREST, F.InterpolationMode.BILINEAR],
    ):
        yield ArgsKwargs(image_loader, angle=15.0, fill=0)
783
784
785


def reference_inputs_rotate_image_tensor():
786
    for image_loader, angle in itertools.product(make_image_loaders_for_interpolation(), _ROTATE_ANGLES):
787
788
789
790
791
792
793
794
        yield ArgsKwargs(image_loader, angle=angle)


def sample_inputs_rotate_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders():
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
795
            spatial_size=bounding_box_loader.spatial_size,
796
797
798
799
            angle=_ROTATE_ANGLES[0],
        )


800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
def reference_inputs_rotate_bounding_box():
    for bounding_box_loader, angle in itertools.product(
        make_bounding_box_loaders(extra_dims=((), (4,))), _ROTATE_ANGLES
    ):
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
            spatial_size=bounding_box_loader.spatial_size,
            angle=angle,
        )

    # TODO: add samples with expand=True and center


def reference_rotate_bounding_box(bounding_box, *, format, spatial_size, angle, expand=False, center=None):

    if center is None:
        center = [spatial_size[1] * 0.5, spatial_size[0] * 0.5]

    a = np.cos(angle * np.pi / 180.0)
    b = np.sin(angle * np.pi / 180.0)
    cx = center[0]
    cy = center[1]
    affine_matrix = np.array(
        [
            [a, b, cx - cx * a - b * cy],
            [-b, a, cy + cx * b - a * cy],
        ],
        dtype="float64" if bounding_box.dtype == torch.float64 else "float32",
    )

    expected_bboxes = reference_affine_bounding_box_helper(
        bounding_box, format=format, spatial_size=spatial_size, affine_matrix=affine_matrix
    )
    return expected_bboxes, spatial_size


837
def sample_inputs_rotate_mask():
838
839
    for mask_loader in make_mask_loaders(sizes=["random"], num_categories=["random"], num_objects=["random"]):
        yield ArgsKwargs(mask_loader, angle=15.0)
840
841


842
843
844
845
846
def sample_inputs_rotate_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, angle=15.0)


847
848
849
850
851
852
853
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.rotate_image_tensor,
            sample_inputs_fn=sample_inputs_rotate_image_tensor,
            reference_fn=pil_reference_wrapper(F.rotate_image_pil),
            reference_inputs_fn=reference_inputs_rotate_image_tensor,
854
            float32_vs_uint8=True,
855
            closeness_kwargs=pil_reference_pixel_difference(1, mae=True),
856
            test_marks=[
857
                xfail_jit_python_scalar_arg("fill"),
858
            ],
859
860
861
862
        ),
        KernelInfo(
            F.rotate_bounding_box,
            sample_inputs_fn=sample_inputs_rotate_bounding_box,
863
864
            reference_fn=reference_rotate_bounding_box,
            reference_inputs_fn=reference_inputs_rotate_bounding_box,
865
            closeness_kwargs={
866
867
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-4, rtol=1e-4),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-4, rtol=1e-4),
868
            },
869
870
871
872
873
        ),
        KernelInfo(
            F.rotate_mask,
            sample_inputs_fn=sample_inputs_rotate_mask,
        ),
874
875
876
877
        KernelInfo(
            F.rotate_video,
            sample_inputs_fn=sample_inputs_rotate_video,
        ),
878
879
880
881
882
883
884
    ]
)

_CROP_PARAMS = combinations_grid(top=[-8, 0, 9], left=[-8, 0, 9], height=[12, 20], width=[12, 20])


def sample_inputs_crop_image_tensor():
885
    for image_loader, params in itertools.product(
886
        make_image_loaders(sizes=[(16, 17)], color_spaces=["RGB"], dtypes=[torch.float32]),
887
888
889
890
891
892
893
894
        [
            dict(top=4, left=3, height=7, width=8),
            dict(top=-1, left=3, height=7, width=8),
            dict(top=4, left=-1, height=7, width=8),
            dict(top=4, left=3, height=17, width=8),
            dict(top=4, left=3, height=7, width=18),
        ],
    ):
895
896
897
898
        yield ArgsKwargs(image_loader, **params)


def reference_inputs_crop_image_tensor():
899
900
901
    for image_loader, params in itertools.product(
        make_image_loaders(extra_dims=[()], dtypes=[torch.uint8]), _CROP_PARAMS
    ):
902
903
904
905
906
907
908
        yield ArgsKwargs(image_loader, **params)


def sample_inputs_crop_bounding_box():
    for bounding_box_loader, params in itertools.product(
        make_bounding_box_loaders(), [_CROP_PARAMS[0], _CROP_PARAMS[-1]]
    ):
909
        yield ArgsKwargs(bounding_box_loader, format=bounding_box_loader.format, **params)
910
911
912


def sample_inputs_crop_mask():
913
914
    for mask_loader in make_mask_loaders(sizes=[(16, 17)], num_categories=["random"], num_objects=["random"]):
        yield ArgsKwargs(mask_loader, top=4, left=3, height=7, width=8)
915
916
917
918
919
920
921


def reference_inputs_crop_mask():
    for mask_loader, params in itertools.product(make_mask_loaders(extra_dims=[()], num_objects=[1]), _CROP_PARAMS):
        yield ArgsKwargs(mask_loader, **params)


922
923
924
925
926
def sample_inputs_crop_video():
    for video_loader in make_video_loaders(sizes=[(16, 17)], num_frames=["random"]):
        yield ArgsKwargs(video_loader, top=4, left=3, height=7, width=8)


927
928
929
930
931
932
def reference_crop_bounding_box(bounding_box, *, format, top, left, height, width):
    affine_matrix = np.array(
        [
            [1, 0, -left],
            [0, 1, -top],
        ],
933
        dtype="float64" if bounding_box.dtype == torch.float64 else "float32",
934
935
    )

936
937
938
939
940
    spatial_size = (height, width)
    expected_bboxes = reference_affine_bounding_box_helper(
        bounding_box, format=format, spatial_size=spatial_size, affine_matrix=affine_matrix
    )
    return expected_bboxes, spatial_size
941
942
943
944
945
946
947
948
949


def reference_inputs_crop_bounding_box():
    for bounding_box_loader, params in itertools.product(
        make_bounding_box_loaders(extra_dims=((), (4,))), [_CROP_PARAMS[0], _CROP_PARAMS[-1]]
    ):
        yield ArgsKwargs(bounding_box_loader, format=bounding_box_loader.format, **params)


950
951
952
953
954
955
956
957
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.crop_image_tensor,
            kernel_name="crop_image_tensor",
            sample_inputs_fn=sample_inputs_crop_image_tensor,
            reference_fn=pil_reference_wrapper(F.crop_image_pil),
            reference_inputs_fn=reference_inputs_crop_image_tensor,
958
            float32_vs_uint8=True,
959
960
961
962
        ),
        KernelInfo(
            F.crop_bounding_box,
            sample_inputs_fn=sample_inputs_crop_bounding_box,
963
964
            reference_fn=reference_crop_bounding_box,
            reference_inputs_fn=reference_inputs_crop_bounding_box,
965
966
967
968
969
970
        ),
        KernelInfo(
            F.crop_mask,
            sample_inputs_fn=sample_inputs_crop_mask,
            reference_fn=pil_reference_wrapper(F.crop_image_pil),
            reference_inputs_fn=reference_inputs_crop_mask,
971
            float32_vs_uint8=True,
972
        ),
973
974
975
976
        KernelInfo(
            F.crop_video,
            sample_inputs_fn=sample_inputs_crop_video,
        ),
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
    ]
)

_RESIZED_CROP_PARAMS = combinations_grid(top=[-8, 9], left=[-8, 9], height=[12], width=[12], size=[(16, 18)])


def sample_inputs_resized_crop_image_tensor():
    for image_loader in make_image_loaders():
        yield ArgsKwargs(image_loader, **_RESIZED_CROP_PARAMS[0])


@pil_reference_wrapper
def reference_resized_crop_image_tensor(*args, **kwargs):
    if not kwargs.pop("antialias", False) and kwargs.get("interpolation", F.InterpolationMode.BILINEAR) in {
        F.InterpolationMode.BILINEAR,
        F.InterpolationMode.BICUBIC,
    }:
        raise pytest.UsageError("Anti-aliasing is always active in PIL")
    return F.resized_crop_image_pil(*args, **kwargs)


def reference_inputs_resized_crop_image_tensor():
    for image_loader, interpolation, params in itertools.product(
1000
        make_image_loaders_for_interpolation(),
1001
1002
        [
            F.InterpolationMode.NEAREST,
1003
            F.InterpolationMode.NEAREST_EXACT,
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
            F.InterpolationMode.BILINEAR,
            F.InterpolationMode.BICUBIC,
        ],
        _RESIZED_CROP_PARAMS,
    ):
        yield ArgsKwargs(
            image_loader,
            interpolation=interpolation,
            antialias=interpolation
            in {
                F.InterpolationMode.BILINEAR,
                F.InterpolationMode.BICUBIC,
            },
            **params,
        )


def sample_inputs_resized_crop_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders():
        yield ArgsKwargs(bounding_box_loader, format=bounding_box_loader.format, **_RESIZED_CROP_PARAMS[0])


def sample_inputs_resized_crop_mask():
    for mask_loader in make_mask_loaders():
        yield ArgsKwargs(mask_loader, **_RESIZED_CROP_PARAMS[0])


1031
1032
1033
1034
1035
def sample_inputs_resized_crop_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, **_RESIZED_CROP_PARAMS[0])


1036
1037
1038
1039
1040
1041
1042
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.resized_crop_image_tensor,
            sample_inputs_fn=sample_inputs_resized_crop_image_tensor,
            reference_fn=reference_resized_crop_image_tensor,
            reference_inputs_fn=reference_inputs_resized_crop_image_tensor,
1043
            float32_vs_uint8=True,
1044
            closeness_kwargs={
1045
                **cuda_vs_cpu_pixel_difference(),
1046
1047
                **pil_reference_pixel_difference(3, mae=True),
                **float32_vs_uint8_pixel_difference(3, mae=True),
1048
            },
1049
1050
1051
1052
1053
1054
1055
1056
1057
        ),
        KernelInfo(
            F.resized_crop_bounding_box,
            sample_inputs_fn=sample_inputs_resized_crop_bounding_box,
        ),
        KernelInfo(
            F.resized_crop_mask,
            sample_inputs_fn=sample_inputs_resized_crop_mask,
        ),
1058
1059
1060
        KernelInfo(
            F.resized_crop_video,
            sample_inputs_fn=sample_inputs_resized_crop_video,
1061
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
1062
        ),
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
    ]
)

_PAD_PARAMS = combinations_grid(
    padding=[[1], [1, 1], [1, 1, 2, 2]],
    padding_mode=["constant", "symmetric", "edge", "reflect"],
)


def sample_inputs_pad_image_tensor():
1073
    make_pad_image_loaders = functools.partial(
1074
        make_image_loaders, sizes=["random"], color_spaces=["RGB"], dtypes=[torch.float32]
1075
1076
1077
1078
1079
1080
1081
1082
1083
    )

    for image_loader, padding in itertools.product(
        make_pad_image_loaders(),
        [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]],
    ):
        yield ArgsKwargs(image_loader, padding=padding)

    for image_loader in make_pad_image_loaders():
1084
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
            yield ArgsKwargs(image_loader, padding=[1], fill=fill)

    for image_loader, padding_mode in itertools.product(
        # We branch for non-constant padding and integer inputs
        make_pad_image_loaders(dtypes=[torch.uint8]),
        ["constant", "symmetric", "edge", "reflect"],
    ):
        yield ArgsKwargs(image_loader, padding=[1], padding_mode=padding_mode)

    # `torch.nn.functional.pad` does not support symmetric padding, and thus we have a custom implementation. Besides
    # negative padding, this is already handled by the inputs above.
    for image_loader in make_pad_image_loaders():
        yield ArgsKwargs(image_loader, padding=[-1], padding_mode="symmetric")
1098
1099
1100


def reference_inputs_pad_image_tensor():
1101
1102
1103
1104
1105
1106
1107
    for image_loader, params in itertools.product(
        make_image_loaders(extra_dims=[()], dtypes=[torch.uint8]), _PAD_PARAMS
    ):
        for fill in get_fills(
            num_channels=image_loader.num_channels,
            dtype=image_loader.dtype,
        ):
1108
1109
1110
1111
            # FIXME: PIL kernel doesn't support sequences of length 1 if the number of channels is larger. Shouldn't it?
            if isinstance(fill, (list, tuple)):
                continue

1112
1113
1114
1115
            yield ArgsKwargs(image_loader, fill=fill, **params)


def sample_inputs_pad_bounding_box():
1116
1117
1118
    for bounding_box_loader, padding in itertools.product(
        make_bounding_box_loaders(), [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]]
    ):
1119
        yield ArgsKwargs(
1120
1121
            bounding_box_loader,
            format=bounding_box_loader.format,
1122
            spatial_size=bounding_box_loader.spatial_size,
1123
1124
            padding=padding,
            padding_mode="constant",
1125
        )
1126
1127
1128


def sample_inputs_pad_mask():
1129
1130
    for mask_loader in make_mask_loaders(sizes=["random"], num_categories=["random"], num_objects=["random"]):
        yield ArgsKwargs(mask_loader, padding=[1])
1131
1132
1133


def reference_inputs_pad_mask():
1134
1135
1136
1137
    for mask_loader, fill, params in itertools.product(
        make_mask_loaders(num_objects=[1], extra_dims=[()]), [None, 127], _PAD_PARAMS
    ):
        yield ArgsKwargs(mask_loader, fill=fill, **params)
1138
1139


1140
1141
1142
1143
1144
def sample_inputs_pad_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, padding=[1])


1145
1146
1147
1148
1149
1150
1151
1152
1153
def reference_pad_bounding_box(bounding_box, *, format, spatial_size, padding, padding_mode):

    left, right, top, bottom = _parse_pad_padding(padding)

    affine_matrix = np.array(
        [
            [1, 0, left],
            [0, 1, top],
        ],
1154
        dtype="float64" if bounding_box.dtype == torch.float64 else "float32",
1155
1156
1157
1158
1159
    )

    height = spatial_size[0] + top + bottom
    width = spatial_size[1] + left + right

1160
1161
1162
    expected_bboxes = reference_affine_bounding_box_helper(
        bounding_box, format=format, spatial_size=(height, width), affine_matrix=affine_matrix
    )
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
    return expected_bboxes, (height, width)


def reference_inputs_pad_bounding_box():
    for bounding_box_loader, padding in itertools.product(
        make_bounding_box_loaders(extra_dims=((), (4,))), [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]]
    ):
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
            spatial_size=bounding_box_loader.spatial_size,
            padding=padding,
            padding_mode="constant",
        )


1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
def pad_xfail_jit_fill_condition(args_kwargs):
    fill = args_kwargs.kwargs.get("fill")
    if not isinstance(fill, (list, tuple)):
        return False
    elif isinstance(fill, tuple):
        return True
    else:  # isinstance(fill, list):
        return all(isinstance(f, int) for f in fill)


1189
1190
1191
1192
1193
1194
1195
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.pad_image_tensor,
            sample_inputs_fn=sample_inputs_pad_image_tensor,
            reference_fn=pil_reference_wrapper(F.pad_image_pil),
            reference_inputs_fn=reference_inputs_pad_image_tensor,
1196
1197
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
1198
            test_marks=[
1199
1200
1201
1202
                xfail_jit_python_scalar_arg("padding"),
                xfail_jit(
                    "F.pad only supports vector fills for list of floats", condition=pad_xfail_jit_fill_condition
                ),
1203
            ],
1204
1205
1206
1207
        ),
        KernelInfo(
            F.pad_bounding_box,
            sample_inputs_fn=sample_inputs_pad_bounding_box,
1208
1209
            reference_fn=reference_pad_bounding_box,
            reference_inputs_fn=reference_inputs_pad_bounding_box,
1210
            test_marks=[
1211
                xfail_jit_python_scalar_arg("padding"),
1212
            ],
1213
1214
1215
1216
1217
1218
        ),
        KernelInfo(
            F.pad_mask,
            sample_inputs_fn=sample_inputs_pad_mask,
            reference_fn=pil_reference_wrapper(F.pad_image_pil),
            reference_inputs_fn=reference_inputs_pad_mask,
1219
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
1220
        ),
1221
1222
1223
1224
        KernelInfo(
            F.pad_video,
            sample_inputs_fn=sample_inputs_pad_video,
        ),
1225
1226
1227
1228
1229
1230
1231
    ]
)

_PERSPECTIVE_COEFFS = [
    [1.2405, 0.1772, -6.9113, 0.0463, 1.251, -5.235, 0.00013, 0.0018],
    [0.7366, -0.11724, 1.45775, -0.15012, 0.73406, 2.6019, -0.0072, -0.0063],
]
1232
1233
_STARTPOINTS = [[0, 1], [2, 3], [4, 5], [6, 7]]
_ENDPOINTS = [[9, 8], [7, 6], [5, 4], [3, 2]]
1234
1235
1236


def sample_inputs_perspective_image_tensor():
1237
    for image_loader in make_image_loaders(sizes=["random"]):
1238
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
1239
1240
1241
1242
1243
            yield ArgsKwargs(
                image_loader, startpoints=None, endpoints=None, fill=fill, coefficients=_PERSPECTIVE_COEFFS[0]
            )

    yield ArgsKwargs(make_image_loader(), startpoints=_STARTPOINTS, endpoints=_ENDPOINTS)
1244
1245
1246


def reference_inputs_perspective_image_tensor():
1247
1248
1249
1250
1251
1252
1253
    for image_loader, coefficients, interpolation in itertools.product(
        make_image_loaders_for_interpolation(),
        _PERSPECTIVE_COEFFS,
        [
            F.InterpolationMode.NEAREST,
            F.InterpolationMode.BILINEAR,
        ],
1254
1255
    ):
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
1256
1257
1258
1259
            # FIXME: PIL kernel doesn't support sequences of length 1 if the number of channels is larger. Shouldn't it?
            if isinstance(fill, (list, tuple)):
                continue

1260
1261
1262
1263
1264
1265
1266
1267
            yield ArgsKwargs(
                image_loader,
                startpoints=None,
                endpoints=None,
                interpolation=interpolation,
                fill=fill,
                coefficients=coefficients,
            )
1268
1269
1270
1271
1272


def sample_inputs_perspective_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders():
        yield ArgsKwargs(
1273
1274
            bounding_box_loader,
            format=bounding_box_loader.format,
1275
            spatial_size=bounding_box_loader.spatial_size,
1276
1277
1278
            startpoints=None,
            endpoints=None,
            coefficients=_PERSPECTIVE_COEFFS[0],
1279
1280
        )

1281
    format = datapoints.BoundingBoxFormat.XYXY
1282
    loader = make_bounding_box_loader(format=format)
1283
    yield ArgsKwargs(
1284
        loader, format=format, spatial_size=loader.spatial_size, startpoints=_STARTPOINTS, endpoints=_ENDPOINTS
1285
1286
    )

1287
1288

def sample_inputs_perspective_mask():
1289
    for mask_loader in make_mask_loaders(sizes=["random"]):
1290
1291
1292
        yield ArgsKwargs(mask_loader, startpoints=None, endpoints=None, coefficients=_PERSPECTIVE_COEFFS[0])

    yield ArgsKwargs(make_detection_mask_loader(), startpoints=_STARTPOINTS, endpoints=_ENDPOINTS)
1293
1294
1295
1296
1297
1298


def reference_inputs_perspective_mask():
    for mask_loader, perspective_coeffs in itertools.product(
        make_mask_loaders(extra_dims=[()], num_objects=[1]), _PERSPECTIVE_COEFFS
    ):
1299
        yield ArgsKwargs(mask_loader, startpoints=None, endpoints=None, coefficients=perspective_coeffs)
1300
1301


1302
1303
def sample_inputs_perspective_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
1304
1305
1306
        yield ArgsKwargs(video_loader, startpoints=None, endpoints=None, coefficients=_PERSPECTIVE_COEFFS[0])

    yield ArgsKwargs(make_video_loader(), startpoints=_STARTPOINTS, endpoints=_ENDPOINTS)
1307
1308


1309
1310
1311
1312
1313
1314
1315
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.perspective_image_tensor,
            sample_inputs_fn=sample_inputs_perspective_image_tensor,
            reference_fn=pil_reference_wrapper(F.perspective_image_pil),
            reference_inputs_fn=reference_inputs_perspective_image_tensor,
1316
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
1317
            closeness_kwargs={
1318
                **pil_reference_pixel_difference(2, mae=True),
1319
1320
                **cuda_vs_cpu_pixel_difference(),
                **float32_vs_uint8_pixel_difference(),
1321
1322
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-5, rtol=1e-5),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-5, rtol=1e-5),
1323
            },
1324
            test_marks=[xfail_jit_python_scalar_arg("fill")],
1325
1326
1327
1328
        ),
        KernelInfo(
            F.perspective_bounding_box,
            sample_inputs_fn=sample_inputs_perspective_bounding_box,
1329
1330
1331
1332
            closeness_kwargs={
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-6, rtol=1e-6),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-6, rtol=1e-6),
            },
1333
1334
1335
1336
1337
1338
        ),
        KernelInfo(
            F.perspective_mask,
            sample_inputs_fn=sample_inputs_perspective_mask,
            reference_fn=pil_reference_wrapper(F.perspective_image_pil),
            reference_inputs_fn=reference_inputs_perspective_mask,
1339
1340
1341
1342
            float32_vs_uint8=True,
            closeness_kwargs={
                (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): dict(atol=10, rtol=0),
            },
1343
1344
1345
1346
        ),
        KernelInfo(
            F.perspective_video,
            sample_inputs_fn=sample_inputs_perspective_video,
1347
1348
            closeness_kwargs={
                **cuda_vs_cpu_pixel_difference(),
1349
1350
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-5, rtol=1e-5),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-5, rtol=1e-5),
1351
            },
1352
1353
1354
1355
1356
        ),
    ]
)


1357
1358
def _get_elastic_displacement(spatial_size):
    return torch.rand(1, *spatial_size, 2)
1359
1360
1361


def sample_inputs_elastic_image_tensor():
1362
    for image_loader in make_image_loaders(sizes=["random"]):
1363
        displacement = _get_elastic_displacement(image_loader.spatial_size)
1364
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
1365
1366
1367
1368
1369
            yield ArgsKwargs(image_loader, displacement=displacement, fill=fill)


def reference_inputs_elastic_image_tensor():
    for image_loader, interpolation in itertools.product(
1370
        make_image_loaders_for_interpolation(),
1371
1372
1373
1374
1375
1376
        [
            F.InterpolationMode.NEAREST,
            F.InterpolationMode.BILINEAR,
            F.InterpolationMode.BICUBIC,
        ],
    ):
1377
        displacement = _get_elastic_displacement(image_loader.spatial_size)
1378
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
1379
1380
1381
1382
1383
            yield ArgsKwargs(image_loader, interpolation=interpolation, displacement=displacement, fill=fill)


def sample_inputs_elastic_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders():
1384
        displacement = _get_elastic_displacement(bounding_box_loader.spatial_size)
1385
1386
1387
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
1388
            spatial_size=bounding_box_loader.spatial_size,
1389
1390
1391
1392
1393
            displacement=displacement,
        )


def sample_inputs_elastic_mask():
1394
    for mask_loader in make_mask_loaders(sizes=["random"]):
1395
1396
1397
1398
        displacement = _get_elastic_displacement(mask_loader.shape[-2:])
        yield ArgsKwargs(mask_loader, displacement=displacement)


1399
1400
1401
1402
1403
1404
def sample_inputs_elastic_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        displacement = _get_elastic_displacement(video_loader.shape[-2:])
        yield ArgsKwargs(video_loader, displacement=displacement)


1405
1406
1407
1408
1409
1410
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.elastic_image_tensor,
            sample_inputs_fn=sample_inputs_elastic_image_tensor,
            reference_inputs_fn=reference_inputs_elastic_image_tensor,
1411
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
1412
            closeness_kwargs={
1413
                **float32_vs_uint8_pixel_difference(6, mae=True),
1414
1415
                **cuda_vs_cpu_pixel_difference(),
            },
1416
            test_marks=[xfail_jit_python_scalar_arg("fill")],
1417
1418
1419
1420
1421
1422
1423
1424
        ),
        KernelInfo(
            F.elastic_bounding_box,
            sample_inputs_fn=sample_inputs_elastic_bounding_box,
        ),
        KernelInfo(
            F.elastic_mask,
            sample_inputs_fn=sample_inputs_elastic_mask,
1425
1426
1427
1428
        ),
        KernelInfo(
            F.elastic_video,
            sample_inputs_fn=sample_inputs_elastic_video,
1429
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
1430
1431
1432
1433
1434
        ),
    ]
)


1435
_CENTER_CROP_SPATIAL_SIZES = [(16, 16), (7, 33), (31, 9)]
1436
_CENTER_CROP_OUTPUT_SIZES = [[4, 3], [42, 70], [4], 3, (5, 2), (6,)]
1437
1438
1439
1440


def sample_inputs_center_crop_image_tensor():
    for image_loader, output_size in itertools.product(
1441
        make_image_loaders(sizes=[(16, 17)], color_spaces=["RGB"], dtypes=[torch.float32]),
1442
1443
1444
1445
1446
1447
        [
            # valid `output_size` types for which cropping is applied to both dimensions
            *[5, (4,), (2, 3), [6], [3, 2]],
            # `output_size`'s for which at least one dimension needs to be padded
            *[[4, 18], [17, 5], [17, 18]],
        ],
1448
1449
1450
1451
1452
1453
    ):
        yield ArgsKwargs(image_loader, output_size=output_size)


def reference_inputs_center_crop_image_tensor():
    for image_loader, output_size in itertools.product(
1454
1455
        make_image_loaders(sizes=_CENTER_CROP_SPATIAL_SIZES, extra_dims=[()], dtypes=[torch.uint8]),
        _CENTER_CROP_OUTPUT_SIZES,
1456
1457
1458
1459
1460
1461
1462
1463
1464
    ):
        yield ArgsKwargs(image_loader, output_size=output_size)


def sample_inputs_center_crop_bounding_box():
    for bounding_box_loader, output_size in itertools.product(make_bounding_box_loaders(), _CENTER_CROP_OUTPUT_SIZES):
        yield ArgsKwargs(
            bounding_box_loader,
            format=bounding_box_loader.format,
1465
            spatial_size=bounding_box_loader.spatial_size,
1466
1467
1468
1469
1470
            output_size=output_size,
        )


def sample_inputs_center_crop_mask():
1471
1472
1473
    for mask_loader in make_mask_loaders(sizes=["random"], num_categories=["random"], num_objects=["random"]):
        height, width = mask_loader.shape[-2:]
        yield ArgsKwargs(mask_loader, output_size=(height // 2, width // 2))
1474
1475
1476
1477


def reference_inputs_center_crop_mask():
    for mask_loader, output_size in itertools.product(
1478
        make_mask_loaders(sizes=_CENTER_CROP_SPATIAL_SIZES, extra_dims=[()], num_objects=[1]), _CENTER_CROP_OUTPUT_SIZES
1479
1480
1481
1482
    ):
        yield ArgsKwargs(mask_loader, output_size=output_size)


1483
1484
1485
1486
1487
1488
def sample_inputs_center_crop_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        height, width = video_loader.shape[-2:]
        yield ArgsKwargs(video_loader, output_size=(height // 2, width // 2))


1489
1490
1491
1492
1493
1494
1495
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.center_crop_image_tensor,
            sample_inputs_fn=sample_inputs_center_crop_image_tensor,
            reference_fn=pil_reference_wrapper(F.center_crop_image_pil),
            reference_inputs_fn=reference_inputs_center_crop_image_tensor,
1496
            float32_vs_uint8=True,
1497
            test_marks=[
1498
                xfail_jit_python_scalar_arg("output_size"),
1499
            ],
1500
1501
1502
1503
        ),
        KernelInfo(
            F.center_crop_bounding_box,
            sample_inputs_fn=sample_inputs_center_crop_bounding_box,
1504
            test_marks=[
1505
                xfail_jit_python_scalar_arg("output_size"),
1506
            ],
1507
1508
1509
1510
1511
1512
        ),
        KernelInfo(
            F.center_crop_mask,
            sample_inputs_fn=sample_inputs_center_crop_mask,
            reference_fn=pil_reference_wrapper(F.center_crop_image_pil),
            reference_inputs_fn=reference_inputs_center_crop_mask,
1513
            float32_vs_uint8=True,
1514
            test_marks=[
1515
                xfail_jit_python_scalar_arg("output_size"),
1516
            ],
1517
        ),
1518
1519
1520
1521
        KernelInfo(
            F.center_crop_video,
            sample_inputs_fn=sample_inputs_center_crop_video,
        ),
1522
1523
1524
1525
1526
    ]
)


def sample_inputs_gaussian_blur_image_tensor():
1527
    make_gaussian_blur_image_loaders = functools.partial(make_image_loaders, sizes=[(7, 33)], color_spaces=["RGB"])
1528
1529
1530
1531
1532
1533

    for image_loader, kernel_size in itertools.product(make_gaussian_blur_image_loaders(), [5, (3, 3), [3, 3]]):
        yield ArgsKwargs(image_loader, kernel_size=kernel_size)

    for image_loader, sigma in itertools.product(
        make_gaussian_blur_image_loaders(), [None, (3.0, 3.0), [2.0, 2.0], 4.0, [1.5], (3.14,)]
1534
    ):
1535
        yield ArgsKwargs(image_loader, kernel_size=5, sigma=sigma)
1536
1537


1538
def sample_inputs_gaussian_blur_video():
1539
    for video_loader in make_video_loaders(sizes=[(7, 33)], num_frames=[5]):
1540
1541
1542
1543
1544
1545
1546
1547
        yield ArgsKwargs(video_loader, kernel_size=[3, 3])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.gaussian_blur_image_tensor,
            sample_inputs_fn=sample_inputs_gaussian_blur_image_tensor,
1548
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
1549
1550
1551
1552
1553
1554
1555
1556
            test_marks=[
                xfail_jit_python_scalar_arg("kernel_size"),
                xfail_jit_python_scalar_arg("sigma"),
            ],
        ),
        KernelInfo(
            F.gaussian_blur_video,
            sample_inputs_fn=sample_inputs_gaussian_blur_video,
1557
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
1558
1559
        ),
    ]
1560
1561
1562
1563
)


def sample_inputs_equalize_image_tensor():
1564
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1565
1566
1567
1568
        yield ArgsKwargs(image_loader)


def reference_inputs_equalize_image_tensor():
1569
1570
1571
    # We are not using `make_image_loaders` here since that uniformly samples the values over the whole value range.
    # Since the whole point of this kernel is to transform an arbitrary distribution of values into a uniform one,
    # the information gain is low if we already provide something really close to the expected value.
1572
    def make_uniform_band_image(shape, dtype, device, *, low_factor, high_factor, memory_format):
1573
1574
1575
1576
1577
1578
1579
        if dtype.is_floating_point:
            low = low_factor
            high = high_factor
        else:
            max_value = torch.iinfo(dtype).max
            low = int(low_factor * max_value)
            high = int(high_factor * max_value)
1580
1581
1582
        return torch.testing.make_tensor(shape, dtype=dtype, device=device, low=low, high=high).to(
            memory_format=memory_format, copy=True
        )
1583

1584
    def make_beta_distributed_image(shape, dtype, device, *, alpha, beta, memory_format):
1585
1586
1587
        image = torch.distributions.Beta(alpha, beta).sample(shape)
        if not dtype.is_floating_point:
            image.mul_(torch.iinfo(dtype).max).round_()
1588
        return image.to(dtype=dtype, device=device, memory_format=memory_format, copy=True)
1589

1590
    spatial_size = (256, 256)
1591
    for dtype, color_space, fn in itertools.product(
1592
        [torch.uint8],
1593
        ["GRAY", "RGB"],
1594
        [
1595
1596
            lambda shape, dtype, device, memory_format: torch.zeros(shape, dtype=dtype, device=device).to(
                memory_format=memory_format, copy=True
1597
            ),
1598
1599
1600
            lambda shape, dtype, device, memory_format: torch.full(
                shape, 1.0 if dtype.is_floating_point else torch.iinfo(dtype).max, dtype=dtype, device=device
            ).to(memory_format=memory_format, copy=True),
1601
            *[
1602
1603
1604
1605
1606
                functools.partial(make_uniform_band_image, low_factor=low_factor, high_factor=high_factor)
                for low_factor, high_factor in [
                    (0.0, 0.25),
                    (0.25, 0.75),
                    (0.75, 1.0),
1607
1608
1609
                ]
            ],
            *[
1610
                functools.partial(make_beta_distributed_image, alpha=alpha, beta=beta)
1611
1612
1613
1614
1615
1616
1617
1618
                for alpha, beta in [
                    (0.5, 0.5),
                    (2, 2),
                    (2, 5),
                    (5, 2),
                ]
            ],
        ],
1619
    ):
1620
        image_loader = ImageLoader(fn, shape=(get_num_channels(color_space), *spatial_size), dtype=dtype)
1621
1622
1623
        yield ArgsKwargs(image_loader)


1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
def sample_inputs_equalize_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader)


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.equalize_image_tensor,
            kernel_name="equalize_image_tensor",
            sample_inputs_fn=sample_inputs_equalize_image_tensor,
            reference_fn=pil_reference_wrapper(F.equalize_image_pil),
1636
            float32_vs_uint8=True,
1637
1638
1639
1640
1641
1642
1643
            reference_inputs_fn=reference_inputs_equalize_image_tensor,
        ),
        KernelInfo(
            F.equalize_video,
            sample_inputs_fn=sample_inputs_equalize_video,
        ),
    ]
1644
1645
1646
1647
)


def sample_inputs_invert_image_tensor():
1648
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1649
1650
1651
1652
        yield ArgsKwargs(image_loader)


def reference_inputs_invert_image_tensor():
1653
    for image_loader in make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]):
1654
1655
1656
        yield ArgsKwargs(image_loader)


1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
def sample_inputs_invert_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader)


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.invert_image_tensor,
            kernel_name="invert_image_tensor",
            sample_inputs_fn=sample_inputs_invert_image_tensor,
            reference_fn=pil_reference_wrapper(F.invert_image_pil),
            reference_inputs_fn=reference_inputs_invert_image_tensor,
1670
            float32_vs_uint8=True,
1671
1672
1673
1674
1675
1676
        ),
        KernelInfo(
            F.invert_video,
            sample_inputs_fn=sample_inputs_invert_video,
        ),
    ]
1677
1678
1679
1680
1681
1682
1683
)


_POSTERIZE_BITS = [1, 4, 8]


def sample_inputs_posterize_image_tensor():
1684
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1685
1686
1687
1688
1689
        yield ArgsKwargs(image_loader, bits=_POSTERIZE_BITS[0])


def reference_inputs_posterize_image_tensor():
    for image_loader, bits in itertools.product(
1690
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1691
1692
1693
1694
1695
        _POSTERIZE_BITS,
    ):
        yield ArgsKwargs(image_loader, bits=bits)


1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
def sample_inputs_posterize_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, bits=_POSTERIZE_BITS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.posterize_image_tensor,
            kernel_name="posterize_image_tensor",
            sample_inputs_fn=sample_inputs_posterize_image_tensor,
            reference_fn=pil_reference_wrapper(F.posterize_image_pil),
            reference_inputs_fn=reference_inputs_posterize_image_tensor,
1709
1710
            float32_vs_uint8=True,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
1711
1712
1713
1714
1715
1716
        ),
        KernelInfo(
            F.posterize_video,
            sample_inputs_fn=sample_inputs_posterize_video,
        ),
    ]
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
)


def _get_solarize_thresholds(dtype):
    for factor in [0.1, 0.5]:
        max_value = get_max_value(dtype)
        yield (float if dtype.is_floating_point else int)(max_value * factor)


def sample_inputs_solarize_image_tensor():
1727
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1728
1729
1730
1731
        yield ArgsKwargs(image_loader, threshold=next(_get_solarize_thresholds(image_loader.dtype)))


def reference_inputs_solarize_image_tensor():
1732
    for image_loader in make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]):
1733
1734
1735
1736
        for threshold in _get_solarize_thresholds(image_loader.dtype):
            yield ArgsKwargs(image_loader, threshold=threshold)


1737
1738
1739
1740
def uint8_to_float32_threshold_adapter(other_args, kwargs):
    return other_args, dict(threshold=kwargs["threshold"] / 255)


1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
def sample_inputs_solarize_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, threshold=next(_get_solarize_thresholds(video_loader.dtype)))


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.solarize_image_tensor,
            kernel_name="solarize_image_tensor",
            sample_inputs_fn=sample_inputs_solarize_image_tensor,
            reference_fn=pil_reference_wrapper(F.solarize_image_pil),
            reference_inputs_fn=reference_inputs_solarize_image_tensor,
1754
1755
            float32_vs_uint8=uint8_to_float32_threshold_adapter,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
1756
1757
1758
1759
1760
1761
        ),
        KernelInfo(
            F.solarize_video,
            sample_inputs_fn=sample_inputs_solarize_video,
        ),
    ]
1762
1763
1764
1765
)


def sample_inputs_autocontrast_image_tensor():
1766
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1767
1768
1769
1770
        yield ArgsKwargs(image_loader)


def reference_inputs_autocontrast_image_tensor():
1771
    for image_loader in make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]):
1772
1773
1774
        yield ArgsKwargs(image_loader)


1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
def sample_inputs_autocontrast_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader)


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.autocontrast_image_tensor,
            kernel_name="autocontrast_image_tensor",
            sample_inputs_fn=sample_inputs_autocontrast_image_tensor,
            reference_fn=pil_reference_wrapper(F.autocontrast_image_pil),
            reference_inputs_fn=reference_inputs_autocontrast_image_tensor,
1788
1789
1790
1791
1792
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(),
            },
1793
1794
1795
1796
1797
1798
        ),
        KernelInfo(
            F.autocontrast_video,
            sample_inputs_fn=sample_inputs_autocontrast_video,
        ),
    ]
1799
1800
1801
1802
1803
1804
1805
1806
)

_ADJUST_SHARPNESS_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_sharpness_image_tensor():
    for image_loader in make_image_loaders(
        sizes=["random", (2, 2)],
1807
        color_spaces=("GRAY", "RGB"),
1808
1809
1810
1811
1812
1813
    ):
        yield ArgsKwargs(image_loader, sharpness_factor=_ADJUST_SHARPNESS_FACTORS[0])


def reference_inputs_adjust_sharpness_image_tensor():
    for image_loader, sharpness_factor in itertools.product(
1814
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1815
1816
1817
1818
1819
        _ADJUST_SHARPNESS_FACTORS,
    ):
        yield ArgsKwargs(image_loader, sharpness_factor=sharpness_factor)


1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
def sample_inputs_adjust_sharpness_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, sharpness_factor=_ADJUST_SHARPNESS_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.adjust_sharpness_image_tensor,
            kernel_name="adjust_sharpness_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_sharpness_image_tensor,
            reference_fn=pil_reference_wrapper(F.adjust_sharpness_image_pil),
            reference_inputs_fn=reference_inputs_adjust_sharpness_image_tensor,
1833
1834
            float32_vs_uint8=True,
            closeness_kwargs=float32_vs_uint8_pixel_difference(2),
1835
1836
1837
1838
1839
1840
        ),
        KernelInfo(
            F.adjust_sharpness_video,
            sample_inputs_fn=sample_inputs_adjust_sharpness_video,
        ),
    ]
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
)


def sample_inputs_erase_image_tensor():
    for image_loader in make_image_loaders(sizes=["random"]):
        # FIXME: make the parameters more diverse
        h, w = 6, 7
        v = torch.rand(image_loader.num_channels, h, w)
        yield ArgsKwargs(image_loader, i=1, j=2, h=h, w=w, v=v)


1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
def sample_inputs_erase_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        # FIXME: make the parameters more diverse
        h, w = 6, 7
        v = torch.rand(video_loader.num_channels, h, w)
        yield ArgsKwargs(video_loader, i=1, j=2, h=h, w=w, v=v)


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.erase_image_tensor,
            kernel_name="erase_image_tensor",
            sample_inputs_fn=sample_inputs_erase_image_tensor,
        ),
        KernelInfo(
            F.erase_video,
            sample_inputs_fn=sample_inputs_erase_video,
        ),
    ]
1872
)
1873
1874
1875
1876
1877

_ADJUST_BRIGHTNESS_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_brightness_image_tensor():
1878
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1879
1880
1881
1882
1883
        yield ArgsKwargs(image_loader, brightness_factor=_ADJUST_BRIGHTNESS_FACTORS[0])


def reference_inputs_adjust_brightness_image_tensor():
    for image_loader, brightness_factor in itertools.product(
1884
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1885
1886
1887
1888
1889
        _ADJUST_BRIGHTNESS_FACTORS,
    ):
        yield ArgsKwargs(image_loader, brightness_factor=brightness_factor)


1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
def sample_inputs_adjust_brightness_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, brightness_factor=_ADJUST_BRIGHTNESS_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.adjust_brightness_image_tensor,
            kernel_name="adjust_brightness_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_brightness_image_tensor,
            reference_fn=pil_reference_wrapper(F.adjust_brightness_image_pil),
            reference_inputs_fn=reference_inputs_adjust_brightness_image_tensor,
1903
1904
            float32_vs_uint8=True,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
1905
1906
1907
1908
1909
1910
        ),
        KernelInfo(
            F.adjust_brightness_video,
            sample_inputs_fn=sample_inputs_adjust_brightness_video,
        ),
    ]
1911
1912
1913
1914
1915
1916
1917
)


_ADJUST_CONTRAST_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_contrast_image_tensor():
1918
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1919
1920
1921
1922
1923
        yield ArgsKwargs(image_loader, contrast_factor=_ADJUST_CONTRAST_FACTORS[0])


def reference_inputs_adjust_contrast_image_tensor():
    for image_loader, contrast_factor in itertools.product(
1924
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1925
1926
1927
1928
1929
        _ADJUST_CONTRAST_FACTORS,
    ):
        yield ArgsKwargs(image_loader, contrast_factor=contrast_factor)


1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
def sample_inputs_adjust_contrast_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, contrast_factor=_ADJUST_CONTRAST_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.adjust_contrast_image_tensor,
            kernel_name="adjust_contrast_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_contrast_image_tensor,
            reference_fn=pil_reference_wrapper(F.adjust_contrast_image_pil),
            reference_inputs_fn=reference_inputs_adjust_contrast_image_tensor,
1943
1944
1945
1946
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(2),
1947
                **cuda_vs_cpu_pixel_difference(),
1948
                (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): pixel_difference_closeness_kwargs(1),
1949
            },
1950
1951
1952
1953
        ),
        KernelInfo(
            F.adjust_contrast_video,
            sample_inputs_fn=sample_inputs_adjust_contrast_video,
1954
1955
1956
1957
            closeness_kwargs={
                **cuda_vs_cpu_pixel_difference(),
                (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): pixel_difference_closeness_kwargs(1),
            },
1958
1959
        ),
    ]
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
)

_ADJUST_GAMMA_GAMMAS_GAINS = [
    (0.5, 2.0),
    (0.0, 1.0),
]


def sample_inputs_adjust_gamma_image_tensor():
    gamma, gain = _ADJUST_GAMMA_GAMMAS_GAINS[0]
1970
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
1971
1972
1973
1974
1975
        yield ArgsKwargs(image_loader, gamma=gamma, gain=gain)


def reference_inputs_adjust_gamma_image_tensor():
    for image_loader, (gamma, gain) in itertools.product(
1976
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1977
1978
1979
1980
1981
        _ADJUST_GAMMA_GAMMAS_GAINS,
    ):
        yield ArgsKwargs(image_loader, gamma=gamma, gain=gain)


1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
def sample_inputs_adjust_gamma_video():
    gamma, gain = _ADJUST_GAMMA_GAMMAS_GAINS[0]
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, gamma=gamma, gain=gain)


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.adjust_gamma_image_tensor,
            kernel_name="adjust_gamma_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_gamma_image_tensor,
            reference_fn=pil_reference_wrapper(F.adjust_gamma_image_pil),
            reference_inputs_fn=reference_inputs_adjust_gamma_image_tensor,
1996
1997
1998
1999
2000
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(),
            },
2001
2002
2003
2004
2005
2006
        ),
        KernelInfo(
            F.adjust_gamma_video,
            sample_inputs_fn=sample_inputs_adjust_gamma_video,
        ),
    ]
2007
2008
2009
2010
2011
2012
2013
)


_ADJUST_HUE_FACTORS = [-0.1, 0.5]


def sample_inputs_adjust_hue_image_tensor():
2014
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
2015
2016
2017
2018
2019
        yield ArgsKwargs(image_loader, hue_factor=_ADJUST_HUE_FACTORS[0])


def reference_inputs_adjust_hue_image_tensor():
    for image_loader, hue_factor in itertools.product(
2020
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
2021
2022
2023
2024
2025
        _ADJUST_HUE_FACTORS,
    ):
        yield ArgsKwargs(image_loader, hue_factor=hue_factor)


2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
def sample_inputs_adjust_hue_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, hue_factor=_ADJUST_HUE_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.adjust_hue_image_tensor,
            kernel_name="adjust_hue_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_hue_image_tensor,
            reference_fn=pil_reference_wrapper(F.adjust_hue_image_pil),
            reference_inputs_fn=reference_inputs_adjust_hue_image_tensor,
2039
2040
            float32_vs_uint8=True,
            closeness_kwargs={
2041
                **pil_reference_pixel_difference(2, mae=True),
2042
2043
                **float32_vs_uint8_pixel_difference(),
            },
2044
2045
2046
2047
2048
2049
        ),
        KernelInfo(
            F.adjust_hue_video,
            sample_inputs_fn=sample_inputs_adjust_hue_video,
        ),
    ]
2050
2051
2052
2053
2054
2055
)

_ADJUST_SATURATION_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_saturation_image_tensor():
2056
    for image_loader in make_image_loaders(sizes=["random"], color_spaces=("GRAY", "RGB")):
2057
2058
2059
2060
2061
        yield ArgsKwargs(image_loader, saturation_factor=_ADJUST_SATURATION_FACTORS[0])


def reference_inputs_adjust_saturation_image_tensor():
    for image_loader, saturation_factor in itertools.product(
2062
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
2063
2064
2065
2066
2067
        _ADJUST_SATURATION_FACTORS,
    ):
        yield ArgsKwargs(image_loader, saturation_factor=saturation_factor)


2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
def sample_inputs_adjust_saturation_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader, saturation_factor=_ADJUST_SATURATION_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.adjust_saturation_image_tensor,
            kernel_name="adjust_saturation_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_saturation_image_tensor,
            reference_fn=pil_reference_wrapper(F.adjust_saturation_image_pil),
            reference_inputs_fn=reference_inputs_adjust_saturation_image_tensor,
2081
2082
2083
2084
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(2),
2085
                **cuda_vs_cpu_pixel_difference(),
2086
            },
2087
2088
2089
2090
        ),
        KernelInfo(
            F.adjust_saturation_video,
            sample_inputs_fn=sample_inputs_adjust_saturation_video,
2091
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
2092
2093
        ),
    ]
2094
2095
2096
2097
2098
2099
)


def sample_inputs_clamp_bounding_box():
    for bounding_box_loader in make_bounding_box_loaders():
        yield ArgsKwargs(
2100
            bounding_box_loader,
2101
2102
            format=bounding_box_loader.format,
            spatial_size=bounding_box_loader.spatial_size,
2103
2104
2105
2106
2107
2108
2109
        )


KERNEL_INFOS.append(
    KernelInfo(
        F.clamp_bounding_box,
        sample_inputs_fn=sample_inputs_clamp_bounding_box,
2110
        logs_usage=True,
2111
2112
2113
2114
2115
2116
    )
)

_FIVE_TEN_CROP_SIZES = [7, (6,), [5], (6, 5), [7, 6]]


2117
def _get_five_ten_crop_spatial_size(size):
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
    if isinstance(size, int):
        crop_height = crop_width = size
    elif len(size) == 1:
        crop_height = crop_width = size[0]
    else:
        crop_height, crop_width = size
    return 2 * crop_height, 2 * crop_width


def sample_inputs_five_crop_image_tensor():
    for size in _FIVE_TEN_CROP_SIZES:
2129
        for image_loader in make_image_loaders(
2130
            sizes=[_get_five_ten_crop_spatial_size(size)],
2131
            color_spaces=["RGB"],
2132
            dtypes=[torch.float32],
2133
        ):
2134
2135
2136
2137
2138
            yield ArgsKwargs(image_loader, size=size)


def reference_inputs_five_crop_image_tensor():
    for size in _FIVE_TEN_CROP_SIZES:
2139
2140
2141
        for image_loader in make_image_loaders(
            sizes=[_get_five_ten_crop_spatial_size(size)], extra_dims=[()], dtypes=[torch.uint8]
        ):
2142
2143
2144
            yield ArgsKwargs(image_loader, size=size)


2145
2146
2147
2148
2149
2150
def sample_inputs_five_crop_video():
    size = _FIVE_TEN_CROP_SIZES[0]
    for video_loader in make_video_loaders(sizes=[_get_five_ten_crop_spatial_size(size)]):
        yield ArgsKwargs(video_loader, size=size)


2151
2152
def sample_inputs_ten_crop_image_tensor():
    for size, vertical_flip in itertools.product(_FIVE_TEN_CROP_SIZES, [False, True]):
2153
        for image_loader in make_image_loaders(
2154
            sizes=[_get_five_ten_crop_spatial_size(size)],
2155
            color_spaces=["RGB"],
2156
            dtypes=[torch.float32],
2157
        ):
2158
2159
2160
2161
2162
            yield ArgsKwargs(image_loader, size=size, vertical_flip=vertical_flip)


def reference_inputs_ten_crop_image_tensor():
    for size, vertical_flip in itertools.product(_FIVE_TEN_CROP_SIZES, [False, True]):
2163
2164
2165
        for image_loader in make_image_loaders(
            sizes=[_get_five_ten_crop_spatial_size(size)], extra_dims=[()], dtypes=[torch.uint8]
        ):
2166
2167
2168
            yield ArgsKwargs(image_loader, size=size, vertical_flip=vertical_flip)


2169
2170
2171
2172
2173
2174
def sample_inputs_ten_crop_video():
    size = _FIVE_TEN_CROP_SIZES[0]
    for video_loader in make_video_loaders(sizes=[_get_five_ten_crop_spatial_size(size)]):
        yield ArgsKwargs(video_loader, size=size)


2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
def multi_crop_pil_reference_wrapper(pil_kernel):
    def wrapper(input_tensor, *other_args, **kwargs):
        output = pil_reference_wrapper(pil_kernel)(input_tensor, *other_args, **kwargs)
        return type(output)(
            F.convert_dtype_image_tensor(F.to_image_tensor(output_pil), dtype=input_tensor.dtype)
            for output_pil in output
        )

    return wrapper


2186
2187
2188
2189
2190
_common_five_ten_crop_marks = [
    xfail_jit_python_scalar_arg("size"),
    mark_framework_limitation(("TestKernels", "test_batched_vs_single"), "Custom batching needed."),
]

2191
2192
2193
2194
2195
KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.five_crop_image_tensor,
            sample_inputs_fn=sample_inputs_five_crop_image_tensor,
2196
            reference_fn=multi_crop_pil_reference_wrapper(F.five_crop_image_pil),
2197
            reference_inputs_fn=reference_inputs_five_crop_image_tensor,
2198
            test_marks=_common_five_ten_crop_marks,
2199
        ),
2200
2201
2202
2203
2204
        KernelInfo(
            F.five_crop_video,
            sample_inputs_fn=sample_inputs_five_crop_video,
            test_marks=_common_five_ten_crop_marks,
        ),
2205
2206
2207
        KernelInfo(
            F.ten_crop_image_tensor,
            sample_inputs_fn=sample_inputs_ten_crop_image_tensor,
2208
            reference_fn=multi_crop_pil_reference_wrapper(F.ten_crop_image_pil),
2209
            reference_inputs_fn=reference_inputs_ten_crop_image_tensor,
2210
            test_marks=_common_five_ten_crop_marks,
2211
        ),
2212
2213
2214
2215
2216
        KernelInfo(
            F.ten_crop_video,
            sample_inputs_fn=sample_inputs_ten_crop_video,
            test_marks=_common_five_ten_crop_marks,
        ),
2217
2218
2219
2220
2221
2222
    ]
)

_NORMALIZE_MEANS_STDS = [
    ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]),
2223
    (0.5, 2.0),
2224
2225
2226
2227
2228
]


def sample_inputs_normalize_image_tensor():
    for image_loader, (mean, std) in itertools.product(
2229
        make_image_loaders(sizes=["random"], color_spaces=["RGB"], dtypes=[torch.float32]),
2230
2231
2232
2233
2234
        _NORMALIZE_MEANS_STDS,
    ):
        yield ArgsKwargs(image_loader, mean=mean, std=std)


2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
def reference_normalize_image_tensor(image, mean, std, inplace=False):
    mean = torch.tensor(mean).view(-1, 1, 1)
    std = torch.tensor(std).view(-1, 1, 1)

    sub = torch.Tensor.sub_ if inplace else torch.Tensor.sub
    return sub(image, mean).div_(std)


def reference_inputs_normalize_image_tensor():
    yield ArgsKwargs(
2245
        make_image_loader(size=(32, 32), color_space="RGB", extra_dims=[1]),
2246
2247
2248
2249
2250
        mean=[0.5, 0.5, 0.5],
        std=[1.0, 1.0, 1.0],
    )


2251
2252
2253
def sample_inputs_normalize_video():
    mean, std = _NORMALIZE_MEANS_STDS[0]
    for video_loader in make_video_loaders(
2254
        sizes=["random"], color_spaces=["RGB"], num_frames=["random"], dtypes=[torch.float32]
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
    ):
        yield ArgsKwargs(video_loader, mean=mean, std=std)


KERNEL_INFOS.extend(
    [
        KernelInfo(
            F.normalize_image_tensor,
            kernel_name="normalize_image_tensor",
            sample_inputs_fn=sample_inputs_normalize_image_tensor,
2265
2266
            reference_fn=reference_normalize_image_tensor,
            reference_inputs_fn=reference_inputs_normalize_image_tensor,
2267
2268
2269
2270
            test_marks=[
                xfail_jit_python_scalar_arg("mean"),
                xfail_jit_python_scalar_arg("std"),
            ],
2271
2272
2273
2274
2275
2276
        ),
        KernelInfo(
            F.normalize_video,
            sample_inputs_fn=sample_inputs_normalize_video,
        ),
    ]
2277
)
2278
2279


2280
def sample_inputs_convert_dtype_image_tensor():
2281
2282
2283
2284
2285
2286
2287
    for input_dtype, output_dtype in itertools.product(
        [torch.uint8, torch.int64, torch.float32, torch.float64], repeat=2
    ):
        if input_dtype.is_floating_point and output_dtype == torch.int64:
            # conversion cannot be performed safely
            continue

2288
        for image_loader in make_image_loaders(sizes=["random"], color_spaces=["RGB"], dtypes=[input_dtype]):
2289
2290
2291
            yield ArgsKwargs(image_loader, dtype=output_dtype)


2292
def reference_convert_dtype_image_tensor(image, dtype=torch.float):
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
    input_dtype = image.dtype
    output_dtype = dtype

    if output_dtype == input_dtype:
        return image

    def fn(value):
        if input_dtype.is_floating_point:
            if output_dtype.is_floating_point:
                return value
            else:
                return int(decimal.Decimal(value) * torch.iinfo(output_dtype).max)
        else:
            input_max_value = torch.iinfo(input_dtype).max

            if output_dtype.is_floating_point:
                return float(decimal.Decimal(value) / input_max_value)
            else:
                output_max_value = torch.iinfo(output_dtype).max

                if input_max_value > output_max_value:
                    factor = (input_max_value + 1) // (output_max_value + 1)
                    return value // factor
                else:
                    factor = (output_max_value + 1) // (input_max_value + 1)
                    return value * factor

    return torch.tensor(tree_map(fn, image.tolist()), dtype=dtype)


2323
def reference_inputs_convert_dtype_image_tensor():
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
    for input_dtype, output_dtype in itertools.product(
        [
            torch.uint8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.float16,
            torch.float32,
            torch.float64,
            torch.bfloat16,
        ],
        repeat=2,
    ):
        if (input_dtype == torch.float32 and output_dtype in {torch.int32, torch.int64}) or (
            input_dtype == torch.float64 and output_dtype == torch.int64
        ):
            continue

        if input_dtype.is_floating_point:
            data = [0.0, 0.5, 1.0]
        else:
            max_value = torch.iinfo(input_dtype).max
            data = [0, max_value // 2, max_value]
        image = torch.tensor(data, dtype=input_dtype)

        yield ArgsKwargs(image, dtype=output_dtype)


2352
2353
2354
2355
2356
def sample_inputs_convert_dtype_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]):
        yield ArgsKwargs(video_loader)


2357
2358
2359
2360
2361
skip_dtype_consistency = TestMark(
    ("TestKernels", "test_dtype_and_device_consistency"),
    pytest.mark.skip(reason="`convert_dtype_*` kernels convert the dtype by design"),
    condition=lambda args_kwargs: args_kwargs.args[0].dtype != args_kwargs.kwargs.get("dtype", torch.float32),
)
2362

2363
2364
2365
KERNEL_INFOS.extend(
    [
        KernelInfo(
2366
2367
2368
2369
            F.convert_dtype_image_tensor,
            sample_inputs_fn=sample_inputs_convert_dtype_image_tensor,
            reference_fn=reference_convert_dtype_image_tensor,
            reference_inputs_fn=reference_inputs_convert_dtype_image_tensor,
2370
            test_marks=[
2371
                skip_dtype_consistency,
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
                TestMark(
                    ("TestKernels", "test_against_reference"),
                    pytest.mark.xfail(reason="Conversion overflows"),
                    condition=lambda args_kwargs: (
                        args_kwargs.args[0].dtype in {torch.float16, torch.bfloat16}
                        and not args_kwargs.kwargs["dtype"].is_floating_point
                    )
                    or (
                        args_kwargs.args[0].dtype in {torch.int32, torch.int64}
                        and args_kwargs.kwargs["dtype"] == torch.float16
                    ),
                ),
            ],
        ),
2386
2387
2388
        KernelInfo(
            F.convert_dtype_video,
            sample_inputs_fn=sample_inputs_convert_dtype_video,
2389
2390
2391
            test_marks=[
                skip_dtype_consistency,
            ],
2392
        ),
2393
2394
    ]
)
2395
2396
2397
2398


def sample_inputs_uniform_temporal_subsample_video():
    for video_loader in make_video_loaders(sizes=["random"], num_frames=[4]):
2399
        yield ArgsKwargs(video_loader, num_samples=2)
2400
2401


2402
def reference_uniform_temporal_subsample_video(x, num_samples):
2403
2404
    # Copy-pasted from
    # https://github.com/facebookresearch/pytorchvideo/blob/c8d23d8b7e597586a9e2d18f6ed31ad8aa379a7a/pytorchvideo/transforms/functional.py#L19
2405
    t = x.shape[-4]
2406
2407
2408
2409
    assert num_samples > 0 and t > 0
    # Sample by nearest neighbor interpolation if num_samples > t.
    indices = torch.linspace(0, t - 1, num_samples)
    indices = torch.clamp(indices, 0, t - 1).long()
2410
    return torch.index_select(x, -4, indices)
2411
2412
2413


def reference_inputs_uniform_temporal_subsample_video():
2414
    for video_loader in make_video_loaders(sizes=["random"], color_spaces=["RGB"], num_frames=[10]):
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
        for num_samples in range(1, video_loader.shape[-4] + 1):
            yield ArgsKwargs(video_loader, num_samples)


KERNEL_INFOS.append(
    KernelInfo(
        F.uniform_temporal_subsample_video,
        sample_inputs_fn=sample_inputs_uniform_temporal_subsample_video,
        reference_fn=reference_uniform_temporal_subsample_video,
        reference_inputs_fn=reference_inputs_uniform_temporal_subsample_video,
    )
)