在Julia中创建一个生成器，它可以同时返回两个值

Question

给定一个生成器：

myVec1 = rand(0:4, 2)
myVec2 = rand(0:4, 8)

myGen = (val1 + val2 for val1 in myVec1, val2 in myVec2)

这基本上是一个有2列的矩阵。可以通过使用collect(myGen)进行查看。

我怎样才能创建一个生成器，每次调用产生两个值（基本上是一列）？

从概念上讲，相当于：

for myCol in eachcol(collect(myGen))
    @show myCol;
end

即使没有显式分配矩阵。

我能否为以下情况包装 myGen：

for value1, value2 in myGen
  dosomethingelse1(value1, value2)
end

换句话说，我想要创建一个生成器，可以返回2个（或更多？）连续的值，并且可以在循环中使用。

所以基本上，我们会在生成器中创建一个二维数组，我想一次访问整个切片。对于实际数组，我可以使用 eachcol 和 eachrow 来完成，但是生成器怎么办呢？

下面是一个测试案例：

myVec1 = rand(0:4, 2);
myVec2 = rand(0:4, 800);

@btime begin
    myMat = [val1 + val2 for val1 in myVec1, val2 in myVec2];
    outVec = [sum(myCol) for myCol in eachcol(myMat)];
end

@btime begin
    myGen = (val1 + val2 for val1 in myVec1, val2 in myVec2);
    outVec = [sum(myCol) for myCol in Iterators.partition(myGen, 2)];
end

@Bogumił Kamiński提供的解决方案确实可行，但实际上，由于某种原因，它会创建更多的分配内存，而不是减少。

- Eric Johnson

Answer 1

这是由Julia语法直接支持的。遍历元组生成器的方式与遍历原子值类似。

例如，您可以尝试：

for (a,b) in ((x, 3x) for x in 1:4)
    println("a=$a, b=$b")
end

- Przemyslaw Szufel

---

我已经更新了问题以进行澄清。不是被调用的函数返回2个值，而是生成器作为2D工作。 - Eric Johnson

你是指通过 IterTools.product 可以实现的笛卡尔积吗？ - Przemyslaw Szufel

Answer 2

我猜您想要像这样的东西：

julia> for (x1, x2) in Iterators.partition(1:10, 2)
           @show x1, x2
       end
(x1, x2) = (1, 2)
(x1, x2) = (3, 4)
(x1, x2) = (5, 6)
(x1, x2) = (7, 8)
(x1, x2) = (9, 10)

如果这是你想要的，那么Iterators.partition是一个可用的函数。

编辑：如果您有两个源流，请使用zip。

julia> for (x1, x2) in zip(1:5, 6:10)
           @show x1, x2
       end
(x1, x2) = (1, 6)
(x1, x2) = (2, 7)
(x1, x2) = (3, 8)
(x1, x2) = (4, 9)
(x1, x2) = (5, 10)

编辑2：我的第一个解决方案已经适用于您的情况。

julia> collect(myGen)
2×8 Matrix{Int64}:
 3  7  5  4  6  3  4  5
 1  5  3  2  4  1  2  3

julia> for (x1, x2) in Iterators.partition(myGen, 2)
           @show x1, x2
       end
(x1, x2) = (3, 1)
(x1, x2) = (7, 5)
(x1, x2) = (5, 3)
(x1, x2) = (4, 2)
(x1, x2) = (6, 4)
(x1, x2) = (3, 1)
(x1, x2) = (4, 2)
(x1, x2) = (5, 3)

- Bogumił Kamiński

---

我更新了我的问题。它还适用吗？ - Eric Johnson

它不起作用了，因为我的第一个解决方案似乎已经按照你的要求完成了。 - Bogumił Kamiński

确实可以工作，但出于某种原因它没有减少分配的数量。有什么想法吗？ - Eric Johnson

问题在于你的 myGen 生成器一开始就不是类型稳定的（你可以通过将代码包装在一个函数中并使用 @code_warntype 进行检查来验证）。如果你想要快速的代码，使用显式的 for 循环，而不要使用生成器。 - Bogumił Kamiński

你链接的评论是正确的。然而，似乎 Iterators.partition 不仅会分配内存（如上面所评论的），而且由于某些原因存在类型不稳定的问题（之前我就看到过，并重新检查后发现它还在存在）。 - Bogumił Kamiński

Answer 3

虽然其他答案在某些方面更加通用，但根据问题编辑中OP增加的信息，更节省内存的选项是使用嵌套生成器。例如:

function solution_nested(v1, v2)
    myGen = ((val1 + val2 for val1 in v1) for val2 in v2)
    [sum(myCol) for myCol in myGen]
end

在测试解决方案时，应避免使用全局变量，并最好将解决方案封装在一个函数中，以便让Julia有足够的机会优化代码。

此解决方案仅产生预期的一个分配结果：

julia> @btime solution_nested(myVec1, myVec2);
  1.856 μs (1 allocation: 6.38 KiB)

虽然这个解决方案与标题不完全匹配，但似乎符合您所描述的。我们使用懒惰列的懒惰序列。 Iterators.partition 之所以慢且内存效率低，是因为它实际上分配了分区中值的中间向量: https://github.com/JuliaLang/julia/blob/dacf9d65aff4668b8fff25957d9aaa2cf03868c8/base/iterators.jl#L1232。

- ahnlabb

---

你的解决方案中唯一的分配是输出的分配吗？ - Eric Johnson

Answer 4

在第二个循环中，您基本上在解构值的元组时缺少括号。更详细地说，您可以在dosomething函数中返回两个值（一个元组）。例如：

function dosomething(element)
    secondElement = element^2
    element, secondElement
end

然后您可以通过解构返回值来使用循环，就像这样：

for (value1, value2) in myGen
    dosomethingelse(value1, value2)
end

如果您需要一个完整的工作示例：

myArray = [1, 2, 3]

function dosomething(element)
    secondElement = element^2
    element, secondElement
end

myGen = (dosomething(myElement) for myElement in myArray)

function dosomethingelse(value1, value2)
    println("Value 1: $value1 \nValue 2: $value2 \n")
end

for (value1, value2) in myGen
    dosomethingelse(value1, value2)
end

---

- Odilf

---

我已经更新了问题以进行澄清。不是被调用的函数返回2个值，而是生成器作为2D工作。 - Eric Johnson