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[, ...])	计算数组分段内的总和。