(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 686 | |
| 687 | @register_lower_rule("PoolingBackwardV1") |
| 688 | def pooling_backward_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 689 | # for forward: y = pool(x) |
| 690 | # for backward: dx = pool_grad(x, y, dy) |
| 691 | assert len(args) == 3 and len(ctx.vars_in) == 3 and len(ctx.vars_out) == 1 |
| 692 | tensor_format, pool_mode = ctx.param["format"], ctx.param["mode"] |
| 693 | |
| 694 | kernel, stride, padding = _get_pool_param( |
| 695 | (ctx.param["window_h"], ctx.param["window_w"]), |
| 696 | (ctx.param["stride_h"], ctx.param["stride_w"]), |
| 697 | (ctx.param["pad_h"], ctx.param["pad_w"]), |
| 698 | tensor_format, |
| 699 | ) |
| 700 | |
| 701 | x, dy = args[0], args[2] |
| 702 | if pool_mode in str(mops.AdaptivePooling.Mode.AVERAGE): |
| 703 | return avgpooling_grad(x, dy, kernel, stride, padding, count_include_pad=True) |
| 704 | elif pool_mode in str(mops.AdaptivePooling.Mode.AVERAGE_COUNT_EXCLUDE_PADDING): |
| 705 | return avgpooling_grad(x, dy, kernel, stride, padding, count_include_pad=False) |
| 706 | else: |
| 707 | assert pool_mode in str( |
| 708 | mops.AdaptivePooling.Mode.MAX |
| 709 | ), f"unknown adaptive pooling mode {pool_mode}" |
| 710 | return maxpooling_grad(x, dy, kernel, stride, padding) |
| 711 | |
| 712 | |
| 713 | def softmax(x: HLOTensor, axis: int = -1): |
nothing calls this directly
no test coverage detected