| 933 | super().__init__() |
| 934 | |
| 935 | def forward( |
| 936 | self, |
| 937 | inp: Tensor, |
| 938 | map_1: Tensor, |
| 939 | map_2: Tensor, |
| 940 | interpolation: str = "linear", |
| 941 | borderMode: str = "replicate", |
| 942 | borderValue: float = 0.0, |
| 943 | ): |
| 944 | assert all( |
| 945 | isinstance(ele, Tensor) for ele in [inp, map_1, map_2] |
| 946 | ), "expected input is megengine.Tensor" |
| 947 | assert map_1.size > 0, "expect map_1.size > 0" |
| 948 | assert map_2.size == 0 or ( |
| 949 | list(map_1.shape) == list(map_2.shape) |
| 950 | ), "expected map_1.size == 0 or map_1.shape == map_2.shape" |
| 951 | map_xy = map_1 |
| 952 | if map_2.size > 0: |
| 953 | ndim = map_1.ndim |
| 954 | map_xy = stack([map_1, map_2], axis=ndim) |
| 955 | out = remap(inp, map_xy, borderMode, borderValue, interpolation) |
| 956 | return out |
| 957 | |
| 958 | |
| 959 | class GaussianBlur(Module): |