Parse pooling op arguments, normalizing scalars to lists. ATen pooling signatures: max_pool{N}d_with_indices(input, kernel_size, stride, padding, dilation, ceil_mode) avg_pool{N}d(input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) Extra args be
(args, ndim, op_name, is_avg_pool=False)
| 3460 | |
| 3461 | |
| 3462 | def _parse_pool_args(args, ndim, op_name, is_avg_pool=False): # noqa: C901 |
| 3463 | """Parse pooling op arguments, normalizing scalars to lists. |
| 3464 | |
| 3465 | ATen pooling signatures: |
| 3466 | max_pool{N}d_with_indices(input, kernel_size, stride, padding, dilation, ceil_mode) |
| 3467 | avg_pool{N}d(input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) |
| 3468 | |
| 3469 | Extra args beyond (input, kernel_size, stride, padding) are accepted only |
| 3470 | when they match safe defaults: |
| 3471 | max_pool: dilation=1, ceil_mode=False |
| 3472 | avg_pool: ceil_mode=False, count_include_pad=True, divisor_override=None |
| 3473 | |
| 3474 | Returns (kernel_size, stride, padding) as lists of length ndim. |
| 3475 | """ |
| 3476 | if is_avg_pool: |
| 3477 | require_args(args, 2, 7, op_name) |
| 3478 | # args[4] = ceil_mode (must be False) |
| 3479 | if len(args) > 4 and args[4]: |
| 3480 | raise ValueError(f"{op_name}: ceil_mode=True is not supported.") |
| 3481 | # args[5] = count_include_pad (must be True) |
| 3482 | if len(args) > 5 and not args[5]: |
| 3483 | raise ValueError(f"{op_name}: count_include_pad=False is not supported.") |
| 3484 | # args[6] = divisor_override (must be None) |
| 3485 | if len(args) > 6 and args[6] is not None: |
| 3486 | raise ValueError(f"{op_name}: divisor_override is not supported.") |
| 3487 | else: |
| 3488 | require_args(args, 2, 6, op_name) |
| 3489 | # args[4] = dilation (must be 1) |
| 3490 | if len(args) > 4: |
| 3491 | dilation = args[4] |
| 3492 | if isinstance(dilation, list): |
| 3493 | if any(d != 1 for d in dilation): |
| 3494 | raise ValueError( |
| 3495 | f"{op_name}: dilation != 1 is not supported, got {dilation}." |
| 3496 | ) |
| 3497 | elif dilation != 1: |
| 3498 | raise ValueError( |
| 3499 | f"{op_name}: dilation != 1 is not supported, got {dilation}." |
| 3500 | ) |
| 3501 | # args[5] = ceil_mode (must be False) |
| 3502 | if len(args) > 5 and args[5]: |
| 3503 | raise ValueError(f"{op_name}: ceil_mode=True is not supported.") |
| 3504 | |
| 3505 | kernel_size = args[1] |
| 3506 | if isinstance(kernel_size, int): |
| 3507 | kernel_size = [kernel_size] * ndim |
| 3508 | |
| 3509 | stride = args[2] if len(args) > 2 and args[2] else kernel_size |
| 3510 | if isinstance(stride, int): |
| 3511 | stride = [stride] * ndim |
| 3512 | if not stride: # empty list means default to kernel_size |
| 3513 | stride = list(kernel_size) |
| 3514 | |
| 3515 | padding = args[3] if len(args) > 3 else [0] * ndim |
| 3516 | if isinstance(padding, int): |
| 3517 | padding = [padding] * ndim |
| 3518 | |
| 3519 | return list(kernel_size), list(stride), list(padding) |
no test coverage detected