Swapping axes

Hey all,

Say I have a rank 4 tensor in the form of NHWC with shape [1, 224, 224, 3]. What is the most elegant and/or efficient way to swap the color channel axes from RGB to BGR without resorting to Python?

Any guidance would be greatly appreciated!


Update: Well I thought I remembered a thread on the mailing list about adding fancy indexing. Sure enough, it’s in the latest toolchain! Here’s the PR. On that end, I’ve gotten something working. I’m still sure there’s a better approach. If anyone knows one I’d be super grateful. Until then this should suffice:

let size = 224
// Random batch of 1
var testBatch = Tensor<Float>(randomUniform: [1, size, size, 3])
// Pull the channels out using TensorRange expression.
let b = testBatch[0, 0..<size, 0..<size, 2]
let g = testBatch[0, 0..<size, 0..<size, 1]
let r = testBatch[0, 0..<size, 0..<size, 0]

// Look at the first 5 elements of the red channel (RGB)
let before = testBatch[0, 0..<size, 0..<5, 0]

// Swap around the channels to BGR order
testBatch[0, 0..<size, 0..<size, 0] = b
testBatch[0, 0..<size, 0..<size, 1] = g
testBatch[0, 0..<size, 0..<size, 2] = r

// Red should be index 2 now, grab the first five and compare...
let after = testBatch[0, 0..<size, 0..<5, 2]

let epsilon = Float(1e-5)
if ((before - after).sum() < epsilon) {
    print("Test Passed!")

Try this trick: https://github.com/BradLarson/AlexNet-Swift/blob/b33e488ea863dd19d412447765c7420a67cb6453/AlexNet/ImageDataset.swift#L68