jax.ops 模块#

已移除在 JAX 0.2.22 版本中已弃用的函数 jax.ops.index_update、jax.ops.index_add 等。请改用 JAX 数组上的 jax.numpy.ndarray.at 属性。

分段归约算子#

segment_max(data, segment_ids[, ...])

计算数组内各分段的最大值。

segment_min(data, segment_ids[, ...])

计算数组内各分段的最小值。

segment_prod(data, segment_ids[, ...])

计算数组内各分段的乘积。

segment_sum(data, segment_ids[, ...])

计算数组内各分段的和。