array_api_extra.at

class array_api_extra.at(x, idx=<object object>, /)

Update operations for read-only arrays.

This implements jax.numpy.ndarray.at for all writeable backends (those that support __setitem__) and routes to the .at[] method for JAX arrays.

Parameters:

  • x (array) – Input array.

  • idx (index, optional) –

    Only array API standard compliant indices are supported.

    You may use two alternate syntaxes:

    >>> import array_api_extra as xpx
>>> xpx.at(x, idx).set(value)  # or add(value), etc.
>>> xpx.at(x)[idx].set(value)

  • copy (bool, optional) –

    None (default)

    The array parameter may be modified in place if it is possible and beneficial for performance. You should not reuse it after calling this function.

    True

    Ensure that the inputs are not modified.

    False

    Ensure that the update operation writes back to the input. Raise ValueError if a copy cannot be avoided.

  • xp (array_namespace, optional) – The standard-compatible namespace for x. Default: infer.

Return type:

Updated input array.

Warning

(a) When you omit the copy parameter, you should never reuse the parameter array later on; ideally, you should reassign it immediately:

>>> import array_api_extra as xpx
>>> x = xpx.at(x, 0).set(2)

The above best practice pattern ensures that the behaviour won’t change depending on whether x is writeable or not, as the original x object is dereferenced as soon as xpx.at returns; this way there is no risk to accidentally update it twice.

On the reverse, the anti-pattern below must be avoided, as it will result in different behaviour on read-only versus writeable arrays:

>>> x = xp.asarray([0, 0, 0])
>>> y = xpx.at(x, 0).set(2)
>>> z = xpx.at(x, 1).set(3)

In the above example, both calls to xpx.at update x in place if possible. This causes the behaviour to diverge depending on whether x is writeable or not:

  • If x is writeable, then after the snippet above you’ll have x == y == z == [2, 3, 0]

  • If x is read-only, then you’ll end up with x == [0, 0, 0], y == [2, 0, 0] and z == [0, 3, 0].

The correct pattern to use if you want diverging outputs from the same input is to enforce copies:

>>> x = xp.asarray([0, 0, 0])
>>> y = xpx.at(x, 0).set(2, copy=True)  # Never updates x
>>> z = xpx.at(x, 1).set(3)  # May or may not update x in place
>>> del x  # avoid accidental reuse of x as we don't know its state anymore

(b) The array API standard does not support integer array indices. The behaviour of update methods when the index is an array of integers is undefined and will vary between backends; this is particularly true when the index contains multiple occurrences of the same index, e.g.:

>>> import numpy as np
>>> import jax.numpy as jnp
>>> import array_api_extra as xpx
>>> xpx.at(np.asarray([123]), np.asarray([0, 0])).add(1)
array([124])
>>> xpx.at(jnp.asarray([123]), jnp.asarray([0, 0])).add(1)
Array([125], dtype=int32)

See also

jax.numpy.ndarray.at

Equivalent array method in JAX.

Notes

sparse, as well as read-only arrays from libraries not explicitly covered by array-api-compat, are not supported by update methods.

Boolean masks are supported on Dask and jitted JAX arrays exclusively when idx has the same shape as x and y is 0-dimensional. Note that this support is not available in JAX’s native x.at[mask].set(y).

This pattern:

>>> mask = m(x)
>>> x[mask] = f(x[mask])

Can’t be replaced by at, as it won’t work on Dask and JAX inside jax.jit:

>>> mask = m(x)
>>> x = xpx.at(x, mask).set(f(x[mask])  # Crash on Dask and jax.jit

You should instead use:

>>> x = xp.where(m(x), f(x), x)

Examples

Given either of these equivalent expressions:

>>> import array_api_extra as xpx
>>> x = xpx.at(x)[1].add(2)
>>> x = xpx.at(x, 1).add(2)

If x is a JAX array, they are the same as:

>>> x = x.at[1].add(2)

If x is a read-only numpy array, they are the same as:

>>> x = x.copy()
>>> x[1] += 2

For other known backends, they are the same as:

>>> x[1] += 2
__init__(x, idx=<object object>, /)
Parameters:

Methods

__init__(x[, idx])

add(y, /[, copy, xp])

Apply x[idx] += y and return the updated array.

divide(y, /[, copy, xp])

Apply x[idx] /= y and return the updated array.

max(y, /[, copy, xp])

Apply x[idx] = maximum(x[idx], y) and return the updated array.

min(y, /[, copy, xp])

Apply x[idx] = minimum(x[idx], y) and return the updated array.

multiply(y, /[, copy, xp])

Apply x[idx] *= y and return the updated array.

power(y, /[, copy, xp])

Apply x[idx] **= y and return the updated array.

set(y, /[, copy, xp])

Apply x[idx] = y and return the update array.

subtract(y, /[, copy, xp])

Apply x[idx] -= y and return the updated array.