jax.ops module

The functions jax.ops.index_update, jax.ops.index_add, etc., which were deprecated in JAX 0.2.22, have been removed. Please use the jax.numpy.ndarray.at property on JAX arrays instead.

Segment reduction operators

segment_max(data, segment_ids[, ...])	Computes the maximum within segments of an array.
segment_min(data, segment_ids[, ...])	Computes the minimum within segments of an array.
segment_prod(data, segment_ids[, ...])	Computes the product within segments of an array.
segment_sum(data, segment_ids[, ...])	Computes the sum within segments of an array.