jax.ops 模块

函数 jax.ops.index_update、jax.ops.index_add 等在 JAX 0.2.22 中已弃用，现已删除。请改用 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[, ...])	计算数组各段中的总和。