jax.vmap

目录

jax.vmap#

jax.vmap(fun, in_axes=0, out_axes=0, axis_name=None, axis_size=None, spmd_axis_name=None)[源代码][源代码]#

向量化映射。创建一个函数，该函数将 fun 映射到参数轴上。

参数:

  • fun (F) – 要在其他轴上进行映射的函数。

  • in_axes (int | None | Sequence[Any]) – 一个整数、None 或值序列，指定要映射的输入数组轴。如果 fun 的每个位置参数都是一个数组，那么 in_axes 可以是一个整数、一个 None 或一个长度等于 fun 位置参数数量的整数和 None 的元组。整数或 None 表示所有参数要映射的数组轴（None 表示不映射任何轴），而元组表示每个相应位置参数要映射的轴。轴整数必须在每个数组的范围 [-ndim, ndim) 内，其中 ndim 是相应输入数组的维度（轴）数。如果 fun 的位置参数是容器（pytree）类型，in_axes 必须是一个序列，长度等于 fun 的位置参数数量，并且对于每个参数，in_axes 中相应的元素可以是一个匹配的 pytree 结构的容器，指定其容器元素的映射。换句话说，in_axes 必须是传递给 fun 的位置参数元组的容器树前缀。更多详情请参见此链接：https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees 必须显式提供 axis_size，或者至少一个位置参数的 in_axes 不能为 None。所有映射位置参数的映射输入轴的大小必须全部相等。作为关键字传递的参数总是映射其前导轴（即轴索引 0）。示例如下。

  • out_axes (Any) – 一个整数、None，或（嵌套的）标准 Python 容器（元组/列表/字典），指示映射轴应出现在输出的哪个位置。所有具有映射轴的输出必须有一个非 None 的 out_axes 规范。轴整数必须在每个输出数组的范围 [-ndim, ndim) 内，其中 ndim 是 vmap() 函数返回的数组的维度（轴）数，比 fun 返回的相应数组的维度（轴）数多一个。

  • axis_name (AxisName | None) – 可选的，一个可哈希的Python对象，用于标识映射的轴，以便可以应用并行集体操作。

  • axis_size (int | None) – 可选，一个表示要映射的轴大小的整数。如果未提供，则映射的轴大小从参数中推断。

  • spmd_axis_name (AxisName | tuple[AxisName, ...] | None)

返回:

批处理/向量化版本的 fun ，其参数对应于 fun 的参数，但在 in_axes 指示的位置有额外的数组轴，并且返回值对应于 fun 的返回值，但在 out_axes 指示的位置有额外的数组轴。

返回类型:

F

例如，我们可以使用向量点积来实现矩阵-矩阵乘积：

>>> import jax.numpy as jnp
>>>
>>> vv = lambda x, y: jnp.vdot(x, y)  #  ([a], [a]) -> []
>>> mv = vmap(vv, (0, None), 0)      #  ([b,a], [a]) -> [b]      (b is the mapped axis)
>>> mm = vmap(mv, (None, 1), 1)      #  ([b,a], [a,c]) -> [b,c]  (c is the mapped axis)

这里我们使用 [a,b] 来表示一个形状为 (a,b) 的数组。以下是一些变体：

>>> mv1 = vmap(vv, (0, 0), 0)   #  ([b,a], [b,a]) -> [b]        (b is the mapped axis)
>>> mv2 = vmap(vv, (0, 1), 0)   #  ([b,a], [a,b]) -> [b]        (b is the mapped axis)
>>> mm2 = vmap(mv2, (1, 1), 0)  #  ([b,c,a], [a,c,b]) -> [c,b]  (c is the mapped axis)

以下是使用 in_axes 中的容器类型来指定映射容器元素的哪些轴的示例：

>>> A, B, C, D = 2, 3, 4, 5
>>> x = jnp.ones((A, B))
>>> y = jnp.ones((B, C))
>>> z = jnp.ones((C, D))
>>> def foo(tree_arg):
...   x, (y, z) = tree_arg
...   return jnp.dot(x, jnp.dot(y, z))
>>> tree = (x, (y, z))
>>> print(foo(tree))
[[12. 12. 12. 12. 12.]
 [12. 12. 12. 12. 12.]]
>>> from jax import vmap
>>> K = 6  # batch size
>>> x = jnp.ones((K, A, B))  # batch axis in different locations
>>> y = jnp.ones((B, K, C))
>>> z = jnp.ones((C, D, K))
>>> tree = (x, (y, z))
>>> vfoo = vmap(foo, in_axes=((0, (1, 2)),))
>>> print(vfoo(tree).shape)
(6, 2, 5)

这里使用 in_axes 中的容器类型（这次是字典）的另一个例子，用于指定要映射的容器元素：

>>> dct = {'a': 0., 'b': jnp.arange(5.)}
>>> x = 1.
>>> def foo(dct, x):
...  return dct['a'] + dct['b'] + x
>>> out = vmap(foo, in_axes=({'a': None, 'b': 0}, None))(dct, x)
>>> print(out)
[1. 2. 3. 4. 5.]

向量化函数的结果可以是映射的或未映射的。例如，下面的函数返回一个对，其中第一个元素是映射的，第二个是未映射的。只有对于未映射的结果，我们可以指定 out_axes 为 ``None``（以保持其未映射状态）。

>>> print(vmap(lambda x, y: (x + y, y * 2.), in_axes=(0, None), out_axes=(0, None))(jnp.arange(2.), 4.))
(Array([4., 5.], dtype=float32), 8.0)

如果为未映射的结果指定了 out_axes ，则结果将在映射的轴上进行广播：

>>> print(vmap(lambda x, y: (x + y, y * 2.), in_axes=(0, None), out_axes=0)(jnp.arange(2.), 4.))
(Array([4., 5.], dtype=float32), Array([8., 8.], dtype=float32, weak_type=True))

如果为映射结果指定了 out_axes ，则结果会相应地进行转置。

最后，这里有一个使用 axis_name 与集体操作的示例：

>>> xs = jnp.arange(3. * 4.).reshape(3, 4)
>>> print(vmap(lambda x: lax.psum(x, 'i'), axis_name='i')(xs))
[[12. 15. 18. 21.]
 [12. 15. 18. 21.]
 [12. 15. 18. 21.]]

参见 jax.pmap() 文档字符串以获取更多涉及集体操作的示例。

目录