CODEWITHSUNDEEP
pythontensorflowdeep-learning
Quail Wwk
Quail Wwk

Reputation: 25

how to drop the lower triangle(include diag) of a 2d tensor in tensorflow1.x?

e.g. I have a tensor:

import tensorflow.compat.v1 as tf
import numpy as np
a = tf.constant(np.array([[1,2,3,4,5],
                          [2,2,4,5,6],
                          [3,4,3,6,7],
                          [4,5,6,4,8],
                          [5,6,7,8,5]))

It's symmetric. Now I only want to see the part where abs(i-j)>s, where i, j denote the row and col index, s is a para.

It equals to j - i >s for symmerty.

So if set s = 2, I want to convert a to:

        tf.constant(np.array([[0,0,0,4,5],
                              [0,0,0,0,6],
                              [0,0,0,0,0],
                              [0,0,0,0,0],
                              [0,0,0,0,0]))

Is there any convince way to do this in tf1.x? TX!

Upvotes: 1

Views: 126

Answers (1)

javidcf
javidcf

Reputation: 59711

You can do that like this:

import tensorflow.compat.v1 as tf
import numpy as np

a = tf.constant(np.array([[1, 2, 3, 4, 5],
                          [2, 2, 4, 5, 6],
                          [3, 4, 3, 6, 7],
                          [4, 5, 6, 4, 8],
                          [5, 6, 7, 8, 5]]))
s = 2
shape = tf.shape(a)
i, j = tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]), indexing='ij')
mask = tf.math.abs(i - j) > s
result = a * tf.dtypes.cast(mask, a.dtype)
tf.print(result)
# [[0 0 0 4 5]
#  [0 0 0 0 6]
#  [0 0 0 0 0]
#  [4 0 0 0 0]
#  [5 6 0 0 0]]

The result is different from what you show, but it is what corresponds to the formula abs(i - j) > s. If you only want the upper part, do instead:

mask = j - i > s

Upvotes: 1

Related Questions