Creating a new matrix from a tensor in tensorflow

Question

Let's say I have a (7,7,3) image with 3 channels. How can I create a matrix A using the image such that the each row consist of just the pairwise connections of the neighbouring pixel? I know this can be done easily in python with for loops but how can we do this within a tensorflow graph?

Example of matrix A:

[[pixel1 pixel2],
 [pixel1 pixel3],
 [pixel1 pixel4],
 .....so on    ]]

Answer 1

You can do this using some matrix algebra. To illustrate the idea, suppose you wanted to do this for a 1D vector.

You can stack the vector with a shifted version of itself to get pairs of neighbors

n = 5
a = tf.range(n)
left = tf.stack([a[1:], a[:n-1]])
left = tf.transpose(left)

By chopping off the tails and repeating for different offset you can get left neighbors and right neighbors

right = tf.stack([a[:n-1], a[1:]])
right = tf.transpose(right)

To ignore edge effects you can chop off the ends and stack again into rank-3 matrix

stacked_neighbors = tf.stack([left[:-1], right[1:]])

Now to interleave the neighbors we can use a trick with transpose and reshape.

stacked_neighbors = tf.transpose(stacked_neighbors, [1, 0, 2])

Since data storage is in row-major order, reshaping into less dimensions than original, reshape flattens excess dimensions on the left

stacked_neighbors = tf.reshape(stacked_neighbors, [6,2])