shuffle a tensor along an axis

Question

shuffle a tensor along an axis

Opened this issue 5 years ago · 8 comments

I'm having trouble shuffling a tensor along a specified axis. I tried using a shuffled array of indices to do this, like so:

var indices: [Int] = Array(0...cardinality)
indices.shuffle()
return data[indices]

The problem is that S4TF does not like integer arrays being used as indices. I get the error argument type '[Int]' does not conform to expected type 'TensorRangeExpression'

Answer 1 · 2020-02-28T06:12:34.000Z

I'm having trouble shuffling a tensor along a specified axis. I tried using a shuffled array of indices to do this, like so:
var indices: [Int] = Array(0...cardinality)
indices.shuffle()
return data[indices]
The problem is that S4TF does not like integer arrays being used as indices. I get the error argument type '[Int]' does not conform to expected type 'TensorRangeExpression'

There's actually a name for this shuffling operation: Tensor.transposed(permutation: [Int]).
_{Edit: transposed(permutation:) is actually different from shuffling! #394 (comment)}

Here's how to use it:

import TensorFlow
let data = Tensor<Float>(randomNormal: [1, 2, 3, 4, 5])
var indices: [Int] = Array(data.indices) // better than `Array(0..<data.rank)`
indices.shuffle()
print(data.shape)
print(data.transposed(permutation: indices).shape)
// [1, 2, 3, 4, 5]
// [4, 1, 5, 3, 2]

Swift supports type extensions, so you can define Tensor.randomPermutation for convenience:

import TensorFlow
extension Tensor {
    var randomPermutation: Tensor {
        var permutation = Array(shape.indices)
        permutation.shuffle()
        return transposed(permutation: permutation)
    }
}
let data = Tensor<Float>(randomNormal: [1, 2, 3, 4, 5])
print(data.shape)
print(data.randomPermutation.shape)
// [1, 2, 3, 4, 5]
// [4, 1, 5, 3, 2]

Hope this helps!

Answer 2 · 2020-02-28T16:14:23.000Z

Thanks for the help! I ran into 2 problems trying to implement this.

First, I'm not entirely sure where the indices attribute is coming from. I'm using the latest toolchain (0.7) and am getting error: value of type 'Tensor<Float>' has no member 'indices' when I run your first example.

Also, it looks like this code would shuffle the order of each axis of a tensor. My goal is to shuffle along an axis. For example:
a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] shuffled along the 0th axis could be [[4, 5, 6], [1, 2, 3], [7, 8, 9]], and shuffled along the 1th axis could be [[2, 1, 3], [5, 4, 6], [8, 7, 9]]. This is like what tf.random.shuffle does

Answer 3 · 2020-02-28T19:58:08.000Z

First, I'm not entirely sure where the indices attribute is coming from. I'm using the latest toolchain (0.7) and am getting error: value of type 'Tensor<Float>' has no member 'indices' when I run your first example.

Sorry! This was a typo that I just fixed. In Tensor.randomPermutation, Array(indices) should be Array(shape.indices).

Also, it looks like this code would shuffle the order of each axis of a tensor. My goal is to shuffle along an axis. For example:
a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] shuffled along the 0th axis could be [[4, 5, 6], [1, 2, 3], [7, 8, 9]], and shuffled along the 1th axis could be [[2, 1, 3], [5, 4, 6], [8, 7, 9]]. This is like what tf.random.shuffle does

Aha, yes. It sounds like you want a "shuffled" function that is different than randomPermutation - sorry for misunderstanding your question (and turning it into an easier one). I think the implementation is a bit more involved, but I have faith you can figure it out!

Answer 4 · 2020-02-28T19:59:37.000Z

Reopening, as the original question (a "shuffle" function like tf.random.shuffle) has not yet been answered.

Answer 5 · 2020-02-28T21:48:51.000Z

This can be done with `_Raw.randomShuffle` if you only shuffle on the first dimension and don’t need the indices anywhere else. Otherwise it can be done with `_Raw.gatherV2` where you specify the axis and provide the indices. I am on mobile but can post examples later. Seems also common enough that adding version outside of _Raw that is differentiable would be handy.

On Sat, Feb 29, 2020 at 5:59 AM Dan Zheng ***@***.***> wrote: Reopened #394 <#394>. — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#394?email_source=notifications&email_token=AAKFCDQOKP3YGYDVNILL5HTRFFUK5A5CNFSM4K4QTJM2YY3PNVWWK3TUL52HS4DFWZEXG43VMVCXMZLOORHG65DJMZUWGYLUNFXW5KTDN5WW2ZLOORPWSZGOW7EAZMQ#event-3083340978>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAKFCDQF44GUVPL4NJVH6DLRFFUK5ANCNFSM4K4QTJMQ> .

-- Michael Kowalski

Answer 6 · 2020-02-29T01:44:42.000Z

An example with _Raw.randomShuffle:

import TensorFlow
let x = Tensor<Int32>(shape: [5, 2], scalars: Array<Int32>(0 ..< 10))
let y = _Raw.randomShuffle(value: x, seed: 1, seed: 2)
print(y) // [[8, 9], [2, 3], [4, 5], [6, 7], [0, 1]]

Even with the seeds I couldn't make this deterministic which surprised me. So the output shows the shuffling but will be different for each run.

An example with _Raw.gatherV2:

import TensorFlow
let x = Tensor<Int32>(shape: [5, 2], scalars: Array<Int32>(0 ..< 10))
let y = _Raw.gatherV2(params: x, indices: Tensor<Int32>([1, 0]), axis: Tensor<Int32>(1))
print(y) // [[1, 0], [3, 2], [5, 4], [7, 6], [9, 8]]

Answer 7 · 2020-03-11T18:15:25.000Z

I will do it.

Answer 8 · 2021-02-16T10:08:22.000Z

here's a quick function that does exactly the same as tf.random.shuffle() but also takes an axis dimension

def tf_shuffle_axis(value, axis=0, seed=None, name=None):
    perm = list(range(tf.rank(value)))
    perm[axis], perm[0] = perm[0], perm[axis]
    value = tf.random.shuffle(tf.transpose(value, perm=perm))
    value = tf.transpose(value, perm=perm)
    return value