jax.numpy.apply_over_axes

jax.numpy.apply_over_axes(func, a, axes)

在指定的轴上重复应用函数。

numpy.apply_over_axes() 的 JAX 实现。

参数:

• func (Callable[[ArrayLike, int], Array]) – 要应用的函数，签名 func(Array, int) -> Array，其中 y = func(x, axis) 必须满足 y.ndim in [x.ndim, x.ndim - 1]。

• a (ArrayLike) – 要在其上应用函数的 N 维数组。

• axes (Sequence[int]) – 要在其上应用函数的轴序列。

返回值:

一个 N 维数组，包含重复函数应用的结果。

返回类型:

Array

另请参阅

• jax.numpy.apply_along_axis()：沿单个轴应用 1D 函数。

示例

此函数旨在具有与典型的关联 jax.numpy 缩减类似的语义，即在一个或多个轴上使用 keepdims=True。

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])

>>> jnp.apply_over_axes(jnp.sum, x, [0])
Array([[5, 7, 9]], dtype=int32)
>>> jnp.sum(x, [0], keepdims=True)
Array([[5, 7, 9]], dtype=int32)

>>> jnp.apply_over_axes(jnp.min, x, [1])
Array([[1],
       [4]], dtype=int32)
>>> jnp.min(x, [1], keepdims=True)
Array([[1],
       [4]], dtype=int32)

>>> jnp.apply_over_axes(jnp.prod, x, [0, 1])
Array([[720]], dtype=int32)
>>> jnp.prod(x, [0, 1], keepdims=True)
Array([[720]], dtype=int32)