jax.ops package¶
Indexed update operators¶
JAX is intended to be used with a functional style of programming, and hence
does not support NumPy-style indexed assignment directly. Instead, JAX provides
pure alternatives, namely jax.ops.index_update()
and its relatives.
Helper object for building indexes for indexed update functions. |
|
|
Pure equivalent of |
|
Pure equivalent of |
|
Pure equivalent of |
|
Pure equivalent of |
|
Pure equivalent of |
Syntactic sugar for indexed update operators¶
JAX also provides an alternate syntax for these indexed update operators.
Specifically, JAX ndarray types have a property at
, which can be used as
follows (where idx
can be an arbitrary index expression).
Alternate syntax |
Equivalent expression |
---|---|
|
|
|
|
|
|
|
|
|
|
Note that none of these expressions modify the original x; instead they return a modified copy of x.
Other operators¶
|
Computes the sum within segments of an array. |