jax.ops package¶

Indexed update operators¶

JAX is intended to be used with a functional style of programming, and does not support NumPy-style indexed assignment directly. Instead, JAX provides alternative pure functional operators for indexed updates to arrays.

JAX array types have a property at, which can be used as follows (where idx is a NumPy index expression).

Alternate syntax

Equivalent in-place expression

x.at[idx].get()

x[idx]

x.at[idx].set(y)

x[idx] = y

x.at[idx].add(y)

x[idx] += y

x.at[idx].multiply(y)

x[idx] *= y

x.at[idx].divide(y)

x[idx] /= y

x.at[idx].power(y)

x[idx] **= y

x.at[idx].min(y)

x[idx] = np.minimum(x[idx], y)

x.at[idx].max(y)

x[idx] = np.maximum(x[idx], y)

None of these expressions modify the original x; instead they return a modified copy of x. However, inside a jit() compiled function, expressions like x = x.at[idx].set(y) are guaranteed to be applied in-place.

By default, JAX assumes that all indices are in-bounds. There is experimental support for giving more precise semantics to out-of-bounds indexed accesses, via the mode parameter to functions such as get and set. Valid values for mode include "clip", which means that out-of-bounds indices will be clamped into range, and "fill"/"drop", which are aliases and mean that out-of-bounds reads will be filled with a scalar fill_value, and out-of-bounds writes will be discarded.

Indexed update functions (deprecated)¶

The following functions are aliases for the x.at[idx].set(y) style operators. Use the x.at[idx] operators instead.

index

Helper object for building indexes for indexed update functions.

index_update(x, idx, y[, …])

Pure equivalent of x[idx] = y.

index_add(x, idx, y[, indices_are_sorted, …])

Pure equivalent of x[idx] += y.

index_mul(x, idx, y[, indices_are_sorted, …])

Pure equivalent of x[idx] *= y.

index_min(x, idx, y[, indices_are_sorted, …])

Pure equivalent of x[idx] = minimum(x[idx], y).

index_max(x, idx, y[, indices_are_sorted, …])

Pure equivalent of x[idx] = maximum(x[idx], y).

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