autopdex.utility.mask_op

autopdex.utility.mask_op(array, selection, values=None, mode='set', ufunc=None)

Performs an operation on a JAX array or dict of arrays based on a boolean mask.

Parameters:

• array (jnp.ndarray or dict) – JAX array or dict of JAX arrays to update.

• selection (jnp.ndarray or dict) – Boolean mask or dict of boolean masks.

• values (jnp.ndarray or dict, optional) – Values to use in the operation. Required for modes except ‘apply’.

• mode (str) – Operation mode. One of ‘set’, ‘add’, ‘multiply’, ‘divide’, ‘power’, ‘min’, ‘max’, ‘apply’.

• ufunc (callable, optional) – A unary function to apply when mode is ‘apply’.

Returns:

Updated array or dict of arrays.