jax.ops module

Contents

`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.