Sorting a matrix by column

user_12504573 0
Viewed 168

I have a matrix which looks like this:

[[5,2],
[4,3],
[3,4]]

Using the command

tf.tensor([...])

And I would like to the matrix by the column with index 0, so that it will look like this:

[[3,4],
[4,3],
[5,2]]

How would I do that using Tensorflow.js?

tensorflow.js tensorflow matrix javascript
3 Answers

Looking at the example, the tensor is sorted over the first axis.

Currently, there is no yet tf.sort in tensorflow.js. But to achieve the same thing, we can slice the tensor over its first axis, then sort it and get the indices. With the later, we can sort the initial tensor using tf.gather

a = tf.tensor2d([[3,4], [4,3],[5,2]])
firstAxis = a.gather([0], 1).reshape([-1]);
ind = tf.topk(firstAxis, a.shape[0]).indices
a.gather(ind.reverse(), 0).print()
user_50699570

You can do it with simple javaScript:

let arr = [[5,2], [4,3], [3,4]];

let result = arr.sort((a, b) => a[0] - b[0]);
console.log(result);

user_136623390

I recently had the same question and used some of edkeveked's answer to come up with these two functions. They look awkward (and probably one can improve on the performance) but they get the job done for a 2D-matrix.

const sortRows = (matrix, asc = true) => {
    return tf.tidy(() => {
        const vector = matrix.reshape([-1])
        const topk = tf.topk(matrix, matrix.shape[1]).indices
        const inds = tf.add(topk, tf.range(0, matrix.shape[0] * matrix.shape[1], matrix.shape[1], 'int32').mul(tf.ones([matrix.shape[1], matrix.shape[0]])).transpose()).reshape([-1]).cast("int32")
        const sorted = vector.gather(inds).reshape(matrix.shape)

        return asc ? sorted.reverse(1) : sorted
    })
}

const sortColumns = (matrix, asc = true) => sortRows(matrix.transpose(), asc).transpose()

Running the following commands

m = tf.tensor([[1, 2, 3], [3, 2, 1], [3, 5, 4], [-2, 1, 6]])
sortColumns(m).print()

will give you

Tensor
    [[-2, 1, 1],
     [1 , 2, 3],
     [3 , 2, 4],
     [3 , 5, 6]]
user_190914890
Related