⚛️
Author mcabbott
Popularity
6 Stars
Updated Last
9 Months Ago
Started In
September 2019

# TransmuteDims.jl

This package provides generalisations of Julia's `permutedims` function and `PermutedDimsArray` wrapper, which allow things other than permutations. These can replace `dropdims` and many uses of `reshape`.

The first generalisation is that you may introduce trivial dimensions. This can be thought of as re-positioning the implicit trivial dimensions beyond `ndims(A)`, such as the 4th and 5th dimensions here:

```A = ones(10,20,30);
ntuple(d -> size(A,d), 5)          # (10, 20, 30, 1, 1)

permutedims(A, (2,3,1)) |> size    # (20, 30, 10)

using TransmuteDims
transmute(A, (4,2,3,5,1)) |> size  # (1, 20, 30, 1, 10)```

Here `(4,2,3,5,1)` is a valid permutation of `1:5`, but the positions of `4,5` don't matter, so in fact this is normalised to `(0,2,3,0,1)`. Zeros indicate trivial output dimensions.

Second, input dimensions below `ndims(A)` may also be omitted, provided they are of size 1:

```A2 = sum(A, dims=2); size(A2)      # (10, 1, 30)
transmute(A2, (3,1)) |> size       # (30, 10)

try transmute(A, (3,1)) catch err; err end  # ArgumentError, "... not allowed when size(A, 2) = 20"```

Third, you may also repeat numbers, to place an input dimension "diagonally" into several output dimensions:

```using LinearAlgebra
transmute(1:10, (1,1)) == Diagonal(1:10)  # true

transmute(A, (2,2,0,3,1)) |> size  # (20, 20, 1, 30, 10)```

The function `transmute` is always lazy, but also tries to minimise the number of wrappers. Ideally to none at all, by un-wrapping and reshaping:

```transmute(A, (4,2,3,5,1)) isa TransmutedDimsArray{Float64, 5, (0,2,3,0,1), (5,2,3), <:Array}

transmute(A, (1,0,2,3)) isa Array{Float64, 4}

transmute(PermutedDimsArray(A, (2,3,1)),(3,1,0,2)) isa Array{Float64, 4}

transmute(Diagonal(1:10), (3,1)) isa TransmutedDimsArray{Int64, 2, (0,1), (2,), <:UnitRange}
transmute(Diagonal(rand(10)), (3,1)) isa Matrix```

Calling the constructor directly `TransmutedDimsArray(A, (3,2,0,1))` simply applies the wrapper. There is also a method `transmute(A, Val((3,2,0,1)))` which works out any un-wrapping at compile-time:

```using BenchmarkTools
@btime transmute(\$A, (2,3,1));           #   6.996 ns (1 allocation: 16 bytes)
@btime PermutedDimsArray(\$A, (2,3,1));   # 386.738 ns (4 allocations: 176 bytes)
@btime transmute(\$A, Val((2,3,1)));      #   1.430 ns (0 allocations: 0 bytes)

@btime transmute(\$A, (1,2,0,3));         #  45.642 ns (2 allocations: 128 bytes)
@btime reshape(\$A, (10,20,1,30));        #  34.479 ns (1 allocation: 80 bytes)```

Finally, there is also an eager variant, which tries always to return a `DenseArray`. This will similarly un-wrap `Transpose` etc, and prefers to reshape if possible, copying data only when necessary. It uses Strided.jl to speed this up, when possible, so should be faster than Base's `permutedims`:

```transmutedims(A, (3,2,0,1)) isa Array{Float64, 4}
transmutedims(1:3, (2,1)) isa Matrix

@btime transmutedims(\$(rand(40,50,60)), (3,2,1));  #  57.365 μs (61 allocations: 944.62 KiB)
@btime permutedims(\$(rand(40,50,60)), (3,2,1));    # 172.643 μs (2 allocations: 937.58 KiB)

@strided(transmute(A, (3,2,0,1))) isa StridedView{Float64, 4}
@strided(transmutedims(A, (3,2,0,1))) isa StridedView{Float64, 4}```

The `StridedView` type is general enough to allow the insertion/removal of trivial dimensions, in addition to permutations, so these functions preserve it.

The lower-case functions also treat tuples as if they were vectors:

```transmute((1,2,3), (1,)) isa AbstractVector
transmutedims((1,2,3), (nothing,1)) isa Matrix```

This was written largely for TensorCast.jl. The immediate issue there was that a `reshape(transpose(::GPUArray))` may fail to trigger GPU broadcasting. This package replaces that with at most one wrapper, ideally none. Calling `transmute` also allowed `@cast` to express what it needs more cleanly.