(ishape, oshape, tensor_format)
| 435 | |
| 436 | |
| 437 | def _get_adaptive_pool_param(ishape, oshape, tensor_format): |
| 438 | assert len(ishape) == 4 and len(oshape) == 4, "only 2-d pooling supported" |
| 439 | if not isinstance(tensor_format, str): |
| 440 | tensor_format = str(tensor_format) |
| 441 | |
| 442 | ishape_hw, oshape_hw = None, None |
| 443 | if tensor_format in str(mops.AdaptivePooling.Format.NCHW): |
| 444 | ishape_hw, oshape_hw = ishape[2:4], oshape[2:4] |
| 445 | elif tensor_format in str(mops.AdaptivePooling.Format.NHWC): |
| 446 | ishape_hw, oshape_hw = ishape[1:3], oshape[1:3] |
| 447 | else: |
| 448 | assert False, f"adaptive pooling only nchw or nhwc, get {tensor_format}" |
| 449 | |
| 450 | assert 0 not in oshape_hw, f"oshape of pooling cannot have zero, get {oshape_hw}" |
| 451 | stride_hw = [(isize // osize) for isize, osize in zip(ishape_hw, oshape_hw)] |
| 452 | kernel_hw = [ |
| 453 | (isize - (osize - 1) * stride) |
| 454 | for isize, osize, stride in zip(ishape_hw, oshape_hw, stride_hw) |
| 455 | ] |
| 456 | |
| 457 | stride, kernel = None, None |
| 458 | if tensor_format in str(mops.AdaptivePooling.Format.NCHW): |
| 459 | stride = (1, 1, *stride_hw) |
| 460 | kernel = (1, 1, *kernel_hw) |
| 461 | elif tensor_format in str(mops.AdaptivePooling.Format.NHWC): |
| 462 | stride = (1, *stride_hw, 1) |
| 463 | kernel = (1, *kernel_hw, 1) |
| 464 | else: |
| 465 | assert False, f"adaptive pooling only nchw or nhwc, get {tensor_format}" |
| 466 | padding = (0, 0, 0, 0) |
| 467 | |
| 468 | return kernel, stride, padding |
| 469 | |
| 470 | |
| 471 | def _select_and_scatter( |
no test coverage detected