I have a tensor A that has shape (batch_size, width, height). Assume that it has these values:
A = torch.tensor([[
                  [0, 1],
                  [1, 0]
]])

I am also given a number K that is a positive integer. Let K=2 in this case. I want to do a process that is similar to reverse pooling and replicate padding. This is the expected output:
B = torch.tensor([[
                [0, 0, 1, 1],
                [0, 0, 1, 1],
                [1, 1, 0, 0],
                [1, 1, 0, 0]             
]])

Explanation: for each element in A, we expand it to the matrix of shape (K, K), and put it in the result tensor. We continue to do this with other elements, and let the stride between them equals to the kernel size (that is, K).
How can I do this in PyTorch? Currently, A is a binary mask, but it could be better if I can expand it to non-binary case.