oryx.core.interpreters.harvest.HarvestTracer

A HarvestTracer just encapsulates a single value.

at Helper property for index update functionality.

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)

aval

Methods

full_lower

View source

get_referent

__abs__

__add__

__and__

__array__

__bool__

__div__

__eq__

Return self==value.

__floordiv__

__ge__

Return self>=value.

__getitem__

__gt__

Return self>value.

__invert__

__iter__

__le__

Return self<=value.

__len__

__lshift__

__lt__

Return self<value.

__matmul__

__mod__

__mul__

__ne__

Return self!=value.

__neg__

__nonzero__

__or__

__pos__

__pow__

__radd__

__rand__

__rdiv__

__rfloordiv__

__rlshift__

__rmatmul__

__rmod__

__rmul__

__ror__

__rpow__

__rrshift__

__rshift__

__rsub__

__rtruediv__

__rxor__

__sub__

__truediv__

__xor__