oryx.core.interpreters.harvest.HarvestTracer

The at property provides a functionally pure equivalent of in-place array modificatons.

============================== ================================ Alternate syntax Equivalent In-place expression ============================== ================================ x = x.at[idx].set(y) x[idx] = y x = x.at[idx].add(y) x[idx] += y x = x.at[idx].multiply(y) x[idx] *= y x = x.at[idx].divide(y) x[idx] /= y x = x.at[idx].power(y) x[idx] **= y x = x.at[idx].min(y) x[idx] = minimum(x[idx], y) x = x.at[idx].max(y) x[idx] = maximum(x[idx], y) x = x.at[idx].apply(ufunc) ufunc.at(x, idx) x = x.at[idx].get() x = x[idx] ============================== ================================

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

Unlike NumPy in-place operations such as :code:x[idx] += y, if multiple indices refer to the same location, all updates will be applied (NumPy would only apply the last update, rather than applying all updates.) The order in which conflicting updates are applied is implementation-defined and may be nondeterministic (e.g., due to concurrency on some hardware platforms).

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 (see below).

Arguments

mode : str Specify out-of-bound indexing mode. Options are:

- ``"promise_in_bounds"``: (default) The user promises that indices are in bounds.
  No additional checking will be performed. In practice, this means that
  out-of-bounds indices in ``get()`` will be clipped, and out-of-bounds indices
  in ``set()``, ``add()``, etc. will be dropped.

- ``"clip"``: clamp out of bounds indices into valid range.
- ``"drop"``: ignore out-of-bound indices.
- ``"fill"``: alias for ``"drop"``.  For `get()`, the optional ``fill_value``
  argument specifies the value that will be returned.

  See :class:`jax.lax.GatherScatterMode` for more details.

indices_are_sorted : bool If True, the implementation will assume that the indices passed to at[] are sorted in ascending order, which can lead to more efficient execution on some backends. unique_indices : bool If True, the implementation will assume that the indices passed to at[] are unique, which can result in more efficient execution on some backends. fill_value : Any Only applies to the get() method: the fill value to return for out-of-bounds slices when mode is 'fill'. Ignored otherwise. Defaults to NaN for inexact types, the largest negative value for signed types, the largest positive value for unsigned types, and True for booleans.

Examples

>>> x = jnp.arange(5.0)
>>> x
DeviceArray([0., 1., 2., 3., 4.], dtype=float32)
>>> x.at[2].add(10)
DeviceArray([ 0.,  1., 12.,  3.,  4.], dtype=float32)
>>> x.at[10].add(10)  # out-of-bounds indices are ignored
DeviceArray([0., 1., 2., 3., 4.], dtype=float32)
>>> x.at[20].add(10, mode='clip')
DeviceArray([ 0.,  1.,  2.,  3., 14.], dtype=float32)
>>> x.at[2].get()
DeviceArray(2., dtype=float32)
>>> x.at[20].get()  # out-of-bounds indices clipped
DeviceArray(4., dtype=float32)
>>> x.at[20].get(mode='fill')  # out-of-bounds indices filled with NaN
DeviceArray(nan, dtype=float32)
>>> x.at[20].get(mode='fill', fill_value=-1)  # custom fill value
DeviceArray(-1., dtype=float32)