Premium problem37. Block-Sum (k×k blocks)

Hard Locked

Given a 2-D tensor Z of shape (h, w) whose dimensions are each divisible by k, compute the sum of each non-overlapping k×k block. Return the reduced tensor of shape (h/k, w/k).

Input

Z =
tensor([[1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        ...,
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.]])
k = 4

Output

tensor([[16., 16., 16., 16.],
        [16., 16., 16., 16.],
        [16., 16., 16., 16.],
        [16., 16., 16., 16.]])

Premium problem

This one's part of Premium. Unlock the full PyTorch track plus every other premium problem on the site.

Implement solve(...)