jax.numpy.apply_over_axes#
- jax.numpy.apply_over_axes(func, a, axes)[source]#
在指定轴上重复应用函数。
JAX 对
numpy.apply_over_axes()
的实现。- 参数:
- 返回:
一个 N 维数组,包含重复应用函数的结果。
- 返回类型:
另请参阅
jax.numpy.apply_along_axis()
:沿单个轴应用一维函数。
示例
此函数旨在与典型的关联性
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)