Create binary tensor from vector in tensorflow

Question

I start with a 1D numpy array x (or tensorflow tensor) with N integer entries. Every entry is smaller or equal to N . I now want to create a tensor Y of shape (N,N) (ie an NxN matrix) where Y[i,j]=0 if x[i]!=x[j] and Y[i,j]=1 if x[i]==x[j] .

Example with numpy:

import numpy as np
x=np.array([1,2,1,2,3,4,2])
Y=np.zeros((x.shape[0],x.shape[0]))
for i in range(x.shape[0]):
    for j in range(x.shape[0]):
        if x[i]==x[j]:
            Y[i,j]=1

Output

array([[ 1.,  0.,  1.,  0.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  1.,  0.,  0.,  1.],
       [ 1.,  0.,  1.,  0.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  1.,  0.,  0.,  1.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  1.,  0.],
       [ 0.,  1.,  0.,  1.,  0.,  0.,  1.]])

How do I create the same function efficiently in pure tensorflow code?

And: What if I have an extra batch dimension, so that the input x has shape (B,N) and I expect as an ouput Y with shape (B,N,N) . The batches are all independent of each other.

Answer 1

Reshape x to two different shapes, (B, 1, N) and (B, N, 1) so they can be properly broadcasted, then compare these two tensors, the result would be what you need with 1 being True and 0 being False :

import tensorflow as tf    
import numpy as np
x=np.array([1,2,1,2,3,4,2])

t = tf.constant(x)

r = tf.cast(
  tf.equal(
    tf.reshape(t, (-1, 1, t.shape[-1].value)), 
    tf.reshape(t, (-1, t.shape[-1].value, 1))
  ), tf.int8)

sess = tf.Session()

sess.run(r)
#array([[[1, 0, 1, 0, 0, 0, 0],
#        [0, 1, 0, 1, 0, 0, 1],
#        [1, 0, 1, 0, 0, 0, 0],
#        [0, 1, 0, 1, 0, 0, 1],
#        [0, 0, 0, 0, 1, 0, 0],
#        [0, 0, 0, 0, 0, 1, 0],
#        [0, 1, 0, 1, 0, 0, 1]]], dtype=int8)

Answer 2

import tensorflow as tf
x = tf.constant([1,2,1,2,3,4,2])
x = tf.expand_dims(x, axis=0)
x = tf.tile(x, [x.shape[1], 1])
x_ = tf.transpose(x)
Y = tf.where(tf.equal(x,x_), tf.ones_like(x), tf.zeros_like(x))

There you have your vector x. You expand dims to have a matrix [1, x.shape]. Then you repeat it to have a copy of the same vector along the lines. x[i] == x[j] is therefore equivalent to x == x_ where x_ is the transposed of your matrix x.

tf.where is a conditional tensor. You give the condition (x == x_), and for each element, if it is true, it will take the first value ( tf.ones_like ) and if it is false it will take the second value ( tf.zeros_like ). Those *_like(x) functions are generating a tensor full of 0 or 1 with the same shape than x.