JAX 中高效规约列表嵌套列表

本文将指导你如何在 JAX 中对嵌套的列表结构进行规约操作，特别是当你需要对多个具有相同结构的列表进行元素级别的求和或类似操作时。 传统的循环方式可能效率较低，而 JAX 提供了更为优雅和高效的解决方案。

JAX 的 jax.tree_util 模块提供了一系列用于处理任意 Python 数据结构的函数，这些数据结构被称为 "PyTrees"。 tree_map 函数允许你将一个函数应用于 PyTree 的每个叶子节点，而 tree_reduce 函数则用于将 PyTree 规约为单个值。

然而，对于列表嵌套列表的规约，直接使用 tree_reduce 可能并不直观。 一个更简洁的方法是结合 tree_map 和 Python 内置的 sum 函数。

使用 tree_map 和 sum 进行规约

假设你有一个包含多个列表的列表 list_of_lists，其中每个子列表具有相同的结构，并且包含 JAX 数组 (jnp.ndarray)。 你的目标是将所有子列表对应位置的元素相加，生成一个新的列表，该列表的结构与子列表相同，但元素是所有对应位置元素之和。

以下是如何使用 tree_map 和 sum 实现此操作的示例代码：

import jax
import jax.numpy as jnp

list_1 = [
    [jnp.asarray([1]), jnp.asarray([2, 3])],
    [jnp.asarray([4]), jnp.asarray([5, 6])],
]

list_2 = [
    [jnp.asarray([7]), jnp.asarray([8, 9])],
    [jnp.asarray([10]), jnp.asarray([11, 12])],
]

list_of_lists = [list_1, list_2]

reduced = jax.tree_util.tree_map(lambda *args: sum(args), *list_of_lists)
print(reduced)

代码解释

表单大师AI

一款基于自然语言处理技术的智能在线表单创建工具，可以帮助用户快速、高效地生成各类专业表单。

74

查看详情

1. *`jax.tree_util.tree_map(function, trees)**:tree_map函数接受一个函数function和一个或多个 PyTrees 作为输入。 在本例中，function是一个 lambda 函数lambda args: sum(args)，而list_of_lists将list_of_lists中的每个子列表作为单独的参数传递给tree_map`。
2. *`lambda args: sum(args)**: 这个 lambda 函数接受任意数量的参数*args，并将它们传递给sum函数。tree_map会遍历所有子列表，并将相同位置的元素作为参数传递给此 lambda 函数。 例如，第一次调用 lambda 函数时，args将包含list_1[0][0]和list_2[0][0]，即jnp.asarray([1])和jnp.asarray([7])`。
3. sum(args): sum 函数将 args 中的所有元素相加。 由于 args 中的元素是 JAX 数组，因此 sum 函数会执行元素级别的加法，并返回一个新的 JAX 数组，其中包含所有输入数组的和。

输出结果

上述代码的输出结果如下：

[[Array([8], dtype=int32), Array([10, 12], dtype=int32)],
 [Array([14], dtype=int32), Array([16, 18], dtype=int32)]]

这正是我们期望的结果：一个新的列表，其结构与原始子列表相同，并且每个元素是所有子列表对应位置元素之和。

注意事项

• tree_map 要求所有输入的 PyTrees 具有相同的结构。 如果子列表的结构不一致，tree_map 将会抛出错误。
• sum 函数适用于 JAX 数组。 如果子列表包含其他类型的元素，你可能需要使用不同的函数来进行规约操作。
• 这种方法可以推广到其他规约操作，例如乘积。 你只需要将 sum 函数替换为相应的函数即可。

总结

通过结合 jax.tree_util.tree_map 和 Python 内置的 sum 函数，你可以高效地对 JAX 中嵌套的列表结构进行规约操作。 这种方法简洁、优雅，并且充分利用了 JAX 的自动微分和编译优化能力。 记住，tree_map 的关键在于确保所有输入的 PyTrees 具有相同的结构，并且选择合适的规约函数来处理叶子节点。

以上就是JAX 中高效规约列表嵌套列表的详细内容，更多请关注php中文网其它相关文章！