jax.numpy.apply_over_axes#

jax.numpy.apply_over_axes(func, a, axes)[原始碼]#

在指定的軸上重複應用函式。

JAX 版本的 numpy.apply_over_axes()。

參數:

  • 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 在一個或多個軸上的縮減操作類似，並使用 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)