shuffle a tensor along an axis

[noahtren]

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'

[dan-zheng]

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! https://github.com/tensorflow/swift/issues/394#issuecomment-592702926}

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!

[noahtren]

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

[dan-zheng]

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!

[dan-zheng]

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

[mikowals]

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 [email protected] wrote:

Reopened #394 https://github.com/tensorflow/swift/issues/394.

— You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub https://github.com/tensorflow/swift/issues/394?email_source=notifications&email_token=AAKFCDQOKP3YGYDVNILL5HTRFFUK5A5CNFSM4K4QTJM2YY3PNVWWK3TUL52HS4DFWZEXG43VMVCXMZLOORHG65DJMZUWGYLUNFXW5KTDN5WW2ZLOORPWSZGOW7EAZMQ#event-3083340978, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAKFCDQF44GUVPL4NJVH6DLRFFUK5ANCNFSM4K4QTJMQ .

-- Michael Kowalski

[mikowals]

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]]

[kiranchhetri1]

I will do it.

[FarisHijazi]

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