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

**URL:** <https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772>\
**Category:** General Usage\
**Tags:** question, arrays, zygote\
**Created:** [September 26, 2021, 6:52am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772 "2021-09-26T06:52:04Z")\
**Posts on this page:** 18\
**Page:** 1

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [September 26, 2021, 6:52am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/1 "2021-09-26T06:52:04Z")

</div>

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,

```julia
[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)

---

<div class="post-metadata">

**Author:** ![oheil](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oheil/32/220745_2.png) [@oheil](https://discourse.julialang.org/u/oheil)\
**Post date:** [September 26, 2021, 8:44am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/2 "2021-09-26T08:44:51Z")

</div>

Welcome!  
Is this what you’re looking for?

```julia
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))

```

---

<div class="post-metadata">

**Author:** ![oheil](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oheil/32/220745_2.png) [@oheil](https://discourse.julialang.org/u/oheil)\
**Post date:** [September 26, 2021, 8:48am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/3 "2021-09-26T08:48:11Z")

</div>

Or how about this:

```julia
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

```

---

<div class="post-metadata">

**Author:** ![aplavin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aplavin/32/222056_2.png) [@aplavin](https://discourse.julialang.org/u/aplavin)\
**Post date:** [September 26, 2021, 9:06am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/4 "2021-09-26T09:06:05Z")

</div>

A go-to solution is the `SplitApplyCombine` package:

```julia
using SplitApplyCombine

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

```

Or use `combinedimsview` for a no-allocation view.

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [September 26, 2021, 1:45pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/5 "2021-09-26T13:45:15Z")

</div>

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!

---

<div class="post-metadata">

**Author:** ![oheil](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oheil/32/220745_2.png) [@oheil](https://discourse.julialang.org/u/oheil)\
**Post date:** [September 26, 2021, 3:44pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/6 "2021-09-26T15:44:47Z")

</div>

> [@tictaccat](#):
>
> the second suggestion of oheil’s doesn’t because I’m trying to write differentiable code, which doesn’t support array mutation

and

> [@tictaccat](#):
>
> a solution that would construct the desired matrix straight away, rather than make an array of arrays, flatten it recursively, and reshape it, I would love to hear it!

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

> [@oheil](#):
>
> `m=zeros(2,3,4,5);`

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.

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [September 26, 2021, 4:03pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/7 "2021-09-26T16:03:40Z")

</div>

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!

---

<div class="post-metadata">

**Author:** ![oheil](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oheil/32/220745_2.png) [@oheil](https://discourse.julialang.org/u/oheil)\
**Post date:** [September 26, 2021, 4:17pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/8 "2021-09-26T16:17:46Z")

</div>

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:

```julia-auto
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

> [@oheil](#):
>
> `m=zeros(2,3,4,5);`

Just for your information.

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [September 26, 2021, 5:14pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/9 "2021-09-26T17:14:28Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![oheil](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oheil/32/220745_2.png) [@oheil](https://discourse.julialang.org/u/oheil)\
**Post date:** [September 26, 2021, 5:18pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/10 "2021-09-26T17:18:05Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [September 26, 2021, 5:47pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/11 "2021-09-26T17:47:05Z")

</div>

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

> [@tictaccat](#):
>
> I want a 2 x 3 x 4 x 5 matrix.

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

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [September 26, 2021, 6:00pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/12 "2021-09-26T18:00:33Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [September 26, 2021, 6:29pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/13 "2021-09-26T18:29:24Z")

</div>

Actually, a matrix is 2D in _every_ language, not just Julia 😉

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

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

```

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [September 30, 2021, 5:16pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/14 "2021-09-30T17:16:09Z")

</div>

> [@DNF](#):
>
> `[f(x,y) for _ in 1:1, _ in 1:1, x in 1:4, y in 1:5]`

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](https://github.com/FluxML/Zygote.jl/issues/421). That basically destroys all my code ☹

---

<div class="post-metadata">

**Author:** ![rafael.guerra](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rafael.guerra/32/216610_2.png) [@rafael.guerra](https://discourse.julialang.org/u/rafael.guerra)\
**Post date:** [September 30, 2021, 10:05pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/15 "2021-09-30T22:05:16Z")

</div>

Note that [TensorCast.jl](https://github.com/mcabbott/TensorCast.jl)’s syntax is quite user friendly for this task:

```julia
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_

```julia
x=1:4; y=1:5; f.(x, y')

```

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [March 20, 2023, 11:35pm UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/16 "2023-03-20T23:35:24Z")

</div>

In Julia 1.9, this is possible with [stack](https://docs.julialang.org/en/v1.9.0-rc1/base/arrays/#Base.stack) 🙂

```julia
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)

```

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [March 21, 2023, 12:29am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/17 "2023-03-21T00:29:09Z")

</div>

Still pre 1.9, this is possible:

```julia
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.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 21, 2023, 12:57am UTC](https://discourse.julialang.org/t/vector-valued-function-in-multidimensional-list-comprehension-differentiable-code/68772/18 "2023-03-21T00:57:21Z")

</div>

> [@Dan](#):
>
> Many solutions don’t avoid allocation

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