Extract boundaries from a tensor with Python

See more:

Python

I have a tensor of size

h=[:,:, 8,8,: ]

The way that I see the tensor is that it is formed by 8 x 8 elements. If we assume periodic boundaries, I want to extract the boundaries of each element,

Let's say that h=[:,:,i,j,:]


I have written the next code:

h=[:,:,i,j,:]

for i in range(8):
    for j in range(8):

        row_start = i-1 if i > 0 else 7
        row_end = i+1 if i < 7 else 0
        col_start = j-1 if j > 0 else 7
        col_end = j+1 if j < 7 else 0


# Extract the submatrix containing the boundaries


        boundaries = h[:,:, row_start:row_end+1, col_start:col_end+1,:]


    # Check if the boundaries extend beyond the edges of the matrix

        if i == 7:
            # Add the top row of the matrix to the top of the boundaries

            top_row = h[:,:,0,  col_start:col_end+1,:]
            top_row = top_row[:,np.newaxis,:,:,:]
            boundaries = torch.vstack([boundaries, top_row])

        if i == 0:

            # Add the bottom row of the matrix to the bottom of the boundaries

            bottom_row = h[:,:,-1,  col_start:col_end+1,:]
            bottom_row = bottom_row[:,np.newaxis,:,:,:]
            boundaries = torch.vstack([boundaries, bottom_row])


        if j == 7:

            # Add the leftmost column of the matrix to the left of the boundaries
            
            left_col = h[:,:,row_start:row_end, 0,:]
            right_col = torch.unsqueeze(right_col, 1)
            boundaries = torch.hstack([boundaries, left_col])   


        if j == 0:

            # Add the rightmost column of the matrix to the right of the boundaries
            right_col = h[:,:,row_start:row_end, -1]
            right_col = torch.unsqueeze(right_col, 1)
            boundaries = torch.hstack([boundaries, right_col])       

What I have tried:

The code is facing problems to extract the boundaries. Any ideas please. Is this code correct?