Vector-valued function in multidimensional list comprehension (Differentiable code)

General Usage
, ,
1

I’ve been using Julia for the first time for a project, and I’ve been running into issues with using list comprehensions to make multidimensional arrays when I have a vector-valued function. As a toy example, I could have a function f(x,y) that outputs a 2 x 3 matrix, and I’d like to run this function on many (x,y) pairs. So far, I’ve been writing list comprehensions like,

[f(x,y) for x in 1:4, y in 1:5]

However, this produces a 4 x 5 matrix of 2D arrays, whereas I want a 2 x 3 x 4 x 5 4D array. I’ve read about some ways to turn the result of this list comprehension into a 4D array, but they all seem quite unwieldy in the case of multidimensional list comprehensions, especially when I’m doing this many times in my code. Is there a clean way to convert this into a 4D array, or even better, a direct method to create the 4D array straight away? Thank you!

(Edit: replaced matrix with 4D array)

2

Welcome!
Is this what you’re looking for?

julia> f(x,y)=fill(x+y,(2,3))
f (generic function with 1 method)

julia> m=[f(x,y) for x in 1:4, y in 1:5]
4×5 Matrix{Matrix{Int64}}:
 [2 2 2; 2 2 2]  [3 3 3; 3 3 3]  [4 4 4; 4 4 4]  [5 5 5; 5 5 5]  [6 6 6; 6 6 6]
 [3 3 3; 3 3 3]  [4 4 4; 4 4 4]  [5 5 5; 5 5 5]  [6 6 6; 6 6 6]  [7 7 7; 7 7 7]
 [4 4 4; 4 4 4]  [5 5 5; 5 5 5]  [6 6 6; 6 6 6]  [7 7 7; 7 7 7]  [8 8 8; 8 8 8]
 [5 5 5; 5 5 5]  [6 6 6; 6 6 6]  [7 7 7; 7 7 7]  [8 8 8; 8 8 8]  [9 9 9; 9 9 9]

julia> r=reshape(collect(Iterators.flatten(m)),(2,3,4,5))
3

Or how about this:

julia> f(x,y)=fill(x+y,(2,3))
f (generic function with 1 method)

julia> m=zeros(2,3,4,5);

julia> for x in 1:4, y in 1:5
       m[:,:,x,y].=f(x,y)
       end
1 Like
4

A go-to solution is the SplitApplyCombine package:

using SplitApplyCombine

[f(x,y) for x in 1:4, y in 1:5] |> combinedims

Or use combinedimsview for a no-allocation view.

5

Thanks for your answers! Both the first suggestion of oheil’s and aplavin’s SplitApplyCombine solution work for me (the second suggestion of oheil’s doesn’t because I’m trying to write differentiable code, which doesn’t support array mutation).

However, both of these suggestions seem to either explicitly or implicitly flatten the array of arrays and then reshape it back up into the desired 4D array. That feels a bit unfortunate because the desired structure is already “baked” into the task itself. I would hope that there’s a way to make the desired 4D array using syntax that is nearly as simple as in the list comprehension. It seems like it would be really easy to run into my problem while working exclusively with multidimensional arrays (and not arrays of arrays). If anyone knows of a solution that would construct the desired 4D array straight away, rather than make an array of arrays, flatten it recursively, and reshape it, I would love to hear it!

6

and

seems contradictory to me.
Can you provide some code (as a minimal working example, MWE) where we can see what you mean with construct? Because I understand this

as the construction.
As the 2x3 matrices are the results of a function, which needs to be called 4x5 times, I can’t see, how this can be done without pre-allocating the target structure, which needs to be mutable therefor.

1 Like
7

I see your point. After thinking about it, it makes sense that I’d only be able to do what I want with pre-allocation and array assignment. Unfortunately, that’s not possible with AD, so I’ll stick to the “vectorized”-style solutions that have been posted in this thread. Thanks for your replies!

8

If your comprehension with subsequent flatten+reshape works, then you have the pre-allocation implicitly in the comprehension itself. If you analyse a simple comprehension you will find:

function vcat(rs::AbstractRange{T}...) where T
    n::Int = 0
    for ra in rs
        n += length(ra)
    end
    a = Vector{T}(undef, n)
    ...
end

where a Vector of length n is allocated and afterwards filled.

From a allocation view point the is no difference between your comprehension and a simple

Just for your information.

2 Likes
9

Right, I see! But it does affect whether Zygote can differentiate it. So I suppose what I was looking for was a single-line function that does the pre-allocation for the desired 4D array implicitly and has an rrule for Zygote to use. It doesn’t look like one exists – I could write one myself, but it seems like this might be hard to do in general (what if f isn’t type stable?). So the two-line solution with flatten+reshape (or just a subsequent list comprehension to collapse the array of arrays into a 4D array) looks like the best bet.

10

I see. Editing the subject to something Zygote/Flux specific may attract some experts on this field.

11

This is confusing. Do you want to create a matrix, or a 4D array?

This is self-contradictory. A matrix is 2D.

12

Sorry for the confusion on terminology, I’m really new to the language. I’d like to create a 4D array.

13

Actually, a matrix is 2D in every language, not just Julia :wink:

Will this work, then? (With f(x,y)=fill(x+y,(2,3))):

jl> [f(x,y) for _ in 1:1, _ in 1:1, x in 1:4, y in 1:5]
1×1×4×5 Array{Matrix{Int64}, 4}:
[:, :, 1, 1] =
 [2 2 2; 2 2 2]

[:, :, 2, 1] =
 [3 3 3; 3 3 3]

[:, :, 3, 1] =
 [4 4 4; 4 4 4]

[:, :, 4, 1] =
 [5 5 5; 5 5 5]

[:, :, 1, 2] =
 [3 3 3; 3 3 3]

[:, :, 2, 2] =
 [4 4 4; 4 4 4]

[:, :, 3, 2] =
 [5 5 5; 5 5 5]

[:, :, 4, 2] =
 [6 6 6; 6 6 6]

[:, :, 1, 3] =
 [4 4 4; 4 4 4]

[:, :, 2, 3] =
 [5 5 5; 5 5 5]

[:, :, 3, 3] =
 [6 6 6; 6 6 6]

[:, :, 4, 3] =
 [7 7 7; 7 7 7]

[:, :, 1, 4] =
 [5 5 5; 5 5 5]

[:, :, 2, 4] =
 [6 6 6; 6 6 6]

[:, :, 3, 4] =
 [7 7 7; 7 7 7]

[:, :, 4, 4] =
 [8 8 8; 8 8 8]

[:, :, 1, 5] =
 [6 6 6; 6 6 6]

[:, :, 2, 5] =
 [7 7 7; 7 7 7]

[:, :, 3, 5] =
 [8 8 8; 8 8 8]

[:, :, 4, 5] =
 [9 9 9; 9 9 9]
1 Like
14

Wow, it turns out that multidimensional list comprehensions aren’t differentiable in Zygote! At least according to No adjoint for Base.Iterators.ProductIterator · Issue #421 · FluxML/Zygote.jl · GitHub. That basically destroys all my code :frowning:

15

Note that TensorCast.jl’s syntax is quite user friendly for this task:

using TensorCast
f(x,y) = fill(x+y,(2,3))
x = 1:4; y = 1:5
@cast m[_,_,i,j] := f(x[i],y[j])  # := returns a view

NB:
your original comprehension could be written simply as

x=1:4;  y=1:5;  f.(x, y')
1 Like
16

In Julia 1.9, this is possible with stack :slightly_smiling_face:

julia> f(x,y) = rand(2,3)
f (generic function with 1 method)

julia> stack([f(x,y) for x in 1:4, y in 1:5]) |> size
(2, 3, 4, 5)
1 Like
17

Still pre 1.9, this is possible:

m = reshape([e for x in 1:4 for y in 1:5 for e in f(x,y)],(2,3,4,5))

Many solutions don’t avoid allocation, as the size of the matrix returned by f is unknown at compile time and it isn’t sure the matrices would fit into a 4D tensor. The above solution works because the comprehensions are not for x, y but of for x for y type giving a vector, which can easily be reshaped.

Perhaps there is a method to avoid allocations using Static Arrays which have compile-time fixed size.

18

If f returns an Array then its allocations will probably dominate. To avoid this, note that besides StaticArrays, stack also accepts tuples. And that it can act on a generator directly, without collecting it:

julia> f2(x,y) = fill(x/y, 2, 3);

julia> @btime stack([f2(x,y) for x in 1:4, y in 1:5]);  # as above
  min 991.600 ns, mean 1.142 μs (23 allocations, 3.50 KiB)

julia> f6(x,y) = ntuple(i -> x/y + i, 6);  # makes a Tuple instead

julia> @btime stack([f6(x,y) for x in 1:4, y in 1:5]);
  min 300.490 ns, mean 350.395 ns (2 allocations, 2.12 KiB)

julia> @btime stack(f6(x,y) for x in 1:4, y in 1:5);  # stack(::Generator) avoids alloc
  min 282.084 ns, mean 300.292 ns (1 allocation, 1.06 KiB)

Not certain these will all be Zygote-friendly right now, but they could be made so.

1 Like