| 401 | |
| 402 | |
| 403 | def fw_conv(input, filter, stride): |
| 404 | # we want to apply 1d conv to any nd array. the way to do it is to reshape |
| 405 | # the input to a 4D tensor. first two dims are singeletons, 3rd dim stores |
| 406 | # all the spatial dims that we are not convolving along now. then we can |
| 407 | # apply conv2d with a 1xK filter. This convolves the same way all the other |
| 408 | # dims stored in the 3d dim. like depthwise conv over these. |
| 409 | # TODO: numpy support |
| 410 | reshaped_input = input.reshape(1, 1, -1, input.shape[-1]) |
| 411 | reshaped_output = torch.nn.functional.conv2d(reshaped_input, |
| 412 | filter.view(1, 1, 1, -1), |
| 413 | stride=(1, stride)) |
| 414 | return reshaped_output.reshape(*input.shape[:-1], -1) |
| 415 | |
| 416 | |
| 417 | def fw_arange(upper_bound, fw, device): |