Tensorflow: Find greater than pairs and stack along axis

Juan Castaño 0
Viewed 10

The problem I have using tensorflow is as follows:

For one tensor X with dims n X m

X = [[x11,x12...,x1m],[x21,x22...,x2m],...[xn1,xn2...,xnm]]

I want to get an n X m X m tensor which are n m X m matrices

Each m X m matrix is the result of:

tf.math.greater(tf.reshape(x,(-1,1)), x) where x is a row  of X

In words, for every row k in X, Im trying to get the pairs i,j where xki > xkj. This gives me a matrix, and then I want to stack those matrices along the first axis, to get a n m x m cube.

Example:

X = [[1,2],[4,3], [5,7]

Result = [[[False, False],[True, False]],[[False, True],[False, False]], [[False, False],[True, False]]]

Result has shape 3 X 2 X 2
python-3.6 tensorflow2.0
1 Answers

Reshaping each row is the same as reshaping all rows. Try this:

def fun(X):
    n, m = X.shape
    X1 = tf.expand_dims(X, -1)
    X2 = tf.reshape(X, (n, 1, m))
    return tf.math.greater(X1, X2)

    
X = tf.Variable([[1,2],[4,3], [5,7]])
print(fun(X))

Output:

tf.Tensor(
[[[False False]
  [ True False]]

 [[False  True]
  [False False]]

 [[False False]
  [ True False]]], shape=(3, 2, 2), dtype=bool)
AndrzejO0
Related