Designating the axes for broadcasting

I can’t find a way to designate the axes over which to broadcast. I looked at broadcast.jl and combine_axes and I can’t make heads or tails of it.

I just want to do

julia> T = randn((3,2,4))
3×2×4 Array{Float64,3}:
[:, :, 1] =
 -1.58198  -1.85354 
  3.75281   0.346631
 -1.58355  -1.78938 

[:, :, 2] =
 -0.995975  -0.590384
 -0.702673   1.73817 
 -1.20793   -1.09226 

[:, :, 3] =
 -0.621339  -1.13472 
  0.671315  -0.81942 
 -0.243524   0.337465

[:, :, 4] =
  1.14836   1.50395 
 -2.23298  -0.819601
  1.19123   0.335359
julia> T2 = randn((3,4))
3×4 Array{Float64,2}:
 -0.88205    0.191643  -0.0997018  -0.720557
  1.01956    0.519003  -1.52737     0.507286
  0.262649  -0.541134  -0.789812   -0.488137

julia> T.*T2
ERROR: DimensionMismatch("arrays could not be broadcast to a common size")
Stacktrace:
 [1] _bcs1 at ./broadcast.jl:438 [inlined]
 [2] _bcs at ./broadcast.jl:432 [inlined] (repeats 2 times)
 [3] broadcast_shape at ./broadcast.jl:426 [inlined]
 [4] combine_axes at ./broadcast.jl:421 [inlined]
 [5] instantiate at ./broadcast.jl:255 [inlined]
 [6] materialize(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{3},Nothing,typeof(*),Tuple{Array{Float64,3},Array{Float64,2}}}) at ./broadcast.jl:753
 [7] top-level scope at none:0

I was hoping there was a simple syntax akin to
broadcast(*,T,T2,dims=(1,3))
I could use permutedims, but that pushes my memory and the actual tensors are considerably larger.

How do I designate dimensions?

Here’s one way:

julia> T3 = T .* reshape(T2, 3,1,4);

julia> size(T3)
(3, 2, 4)

You can also do this, at the cost of copying T2:

julia> T4 = T .* T2[:,[CartesianIndex()],:];

julia> T4 == T3
true

Here’s another way, the function orient just does the same reshape, but is smart enough to copy instead if say T2::Transpose which will be very slow once reshaped:

julia> using TensorCast

julia> T3 == @cast T5[c,b,d] := T[c,b,d] * T2[c,d]
true

julia> @pretty @cast T5[c,b,d] := T[c,b,d] * T2[c,d]
begin
    local penguin = orient(T2, (:, *, :))
    T5 = @__dot__(T * penguin)
end

One more:

julia> T6 = PermutedDimsArray(T, (1,3,2)) .* T2;

julia> size(T6)
(3, 4, 2)
2 Likes

The PermutedDimsArray does the trick for me. Thanks!