input_p = input[:, :, i + 1]
gt_mask = input_p > max_val
max_val = max_val.masked_scatter(gt_mask, input_p[gt_mask])
max_ind.masked_fill_(gt_mask, i + 1)
grad_out_p = grad_out[:, :, i + 1].unsqueeze(2)
output.scatter_add_(2, max_ind.unsqueeze(2), grad_out_p)
return output
def right_pool_forward(input):
pass
def left_pool_forward(input):
pass
if __name__ == "__main__":
import top_pool
import bottom_pool
f = torch.rand(1, 1, 3, 3)
g = torch.rand(1, 1, 3, 3)
print(f, '\n', g)
print(top_pool_backward(f, g))
print(bottom_pool_backward(f, g))
print(top_pool.backward(f, g))
print(bottom_pool.backward(f, g))
def backward(ctx, grad_output):
input = ctx.saved_variables[0]
output = bottom_pool.backward(input, grad_output)[0]
return output