CUBE tensor get 6 face tensor then convert into (1,6xN)

Is there a convenient way to get the values of the six sides of the cube tensor, such as the tensor of NxNxN, get the tensor of NxN with 6 sides, and form the 1d tensor of N values from the largest value of each row in NxN?

Could you write a (slow) reference code which we could use for a proper solution to compare against, please?

Certainly! Here’s a reference code( from chatgpt) that provides a slower, but straightforward solution for extracting the values of the six sides of a cube tensor and forming a 1D tensor with the largest value from each row in the extracted NxN tensors:

import torch

# Assume you have an NxNxN cube tensor named cube_tensor
N = cube_tensor.size(0)  # Get the size of the cube tensor

# Extract each side (NxN tensor) using slicing
front_side = cube_tensor[:, :, 0]
back_side = cube_tensor[:, :, -1]
top_side = cube_tensor[:, 0, :]
bottom_side = cube_tensor[:, -1, :]
left_side = cube_tensor[0, :, :]
right_side = cube_tensor[-1, :, :]

# Form a 1D tensor with the largest value from each row in each side
front_largest_per_row = torch.max(front_side, dim=1).values
back_largest_per_row = torch.max(back_side, dim=1).values
top_largest_per_row = torch.max(top_side, dim=1).values
bottom_largest_per_row = torch.max(bottom_side, dim=1).values
left_largest_per_row = torch.max(left_side, dim=1).values
right_largest_per_row = torch.max(right_side, dim=1).values

# Concatenate the largest values from each side into a single 1D tensor
largest_values =[front_largest_per_row, back_largest_per_row,
                            top_largest_per_row, bottom_largest_per_row,
                            left_largest_per_row, right_largest_per_row])

Is there a more concise way to implement? @ptrblck