sh1ng
sh1ng

Reputation: 2973

Binary matrix, where 1 is located at the maximum value along the row

Let say I have a tensor

[[0.3, 0.7],
[0.9,  0.1]]

How can I create a tensor with 1.0 at maximum positions along axis, so the result should be for axis=1

[[0., 1.],
[1.,  0.]]

In my case first dimension is a batch size, so it's '?'

Upvotes: 0

Views: 212

Answers (2)

Salvador Dali
Salvador Dali

Reputation: 222541

Both of the answers presented are inefficient in terms of memory/compute.

You can calculate it in linear time (no-matmul) without allocating unnecessary memory in just one line:

tf.cast(tf.equal(a, tf.reshape(tf.reduce_max(a, axis=1), (-1, 1))), tf.int16)

The full example is here:

import tensorflow as tf
a = tf.constant([
    [1, 9, 1, 6],
    [6, 5, 0, 6],
    [4, 0, 7, 6],
    [1, 5, 9, 1]
])
b = tf.cast(tf.equal(a, tf.reshape(tf.reduce_max(a, axis=1), (-1, 1))), tf.int16)

with tf.Session() as sess:
    print sess.run(b)

Which will give you

[[0 1 0 0]
 [1 0 0 1]
 [0 0 1 0]
 [0 0 1 0]]

As you see it uses broadcasting in tf.equal to reduce the number of memory.

Upvotes: 1

Zehao Shi
Zehao Shi

Reputation: 99

You can use tf.arg_max() to index the maximum value of the matrix, then use tf.sparse_to_dense() to derive a new matrix you want.

Here is an example.

import tensorflow as tf

a = tf.convert_to_tensor([[0.3, 0.7],
                          [0.9, 0.1]], dtype=tf.float32)

clo_idx = tf.arg_max(a, dimension=0)
row_idx = tf.range(a.get_shape()[0], dtype=tf.int64)

index = tf.concat([tf.reshape(row_idx, shape=[-1, 1]), 
                   tf.reshape(clo_idx, shape=[-1, 1])], axis=1)

b = tf.sparse_to_dense(index, 
                       output_shape=tf.shape(a, out_type=tf.int64),
                       sparse_values=1., 
                       default_value=0.)

with tf.Session() as sess:
  print sess.run(b)

'b' is what you want. And you can get other forms of matrices by setting 'output_shape', 'sparse_values' and 'default_value'.

Upvotes: 0

Related Questions