(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 588 | |
| 589 | @register_lower_rule(mops.AdaptivePooling) |
| 590 | def adaptive_pooling_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 591 | assert len(ctx.vars_in) == 2 and len(args) == 2 and len(ctx.vars_out) == 1 |
| 592 | assert ctx.op.shape == ctx.vars_in[1].bound_data.tolist() and len(ctx.op.shape) == 2 |
| 593 | |
| 594 | ishape, oshape = ctx.vars_in[0].shape, ctx.vars_out[0].shape |
| 595 | kernel, stride, padding = _get_adaptive_pool_param(ishape, oshape, ctx.op.format) |
| 596 | |
| 597 | if ctx.op.mode == mops.AdaptivePooling.Mode.AVERAGE: |
| 598 | return avgpooling( |
| 599 | args[0], stride, kernel, padding, count_include_pad=True, oshape=oshape |
| 600 | ) |
| 601 | elif ctx.op.mode == mops.AdaptivePooling.Mode.AVERAGE_COUNT_EXCLUDE_PADDING: |
| 602 | return avgpooling( |
| 603 | args[0], stride, kernel, padding, count_include_pad=False, oshape=oshape |
| 604 | ) |
| 605 | else: |
| 606 | assert ( |
| 607 | ctx.op.mode == mops.AdaptivePooling.Mode.MAX |
| 608 | ), f"unknown adaptive pooling mode {ctx.op.mode}" |
| 609 | return maxpooling(args[0], stride, kernel, padding, oshape=oshape) |
| 610 | |
| 611 | |
| 612 | @register_lower_rule("AdaptivePoolingBackwardV1") |
nothing calls this directly
no test coverage detected