autopdex.utility.mask_set

autopdex.utility.mask_set(array, selection, values)

Changes values in 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) – Values to insert, or dict of values.

Returns:

Updated array or dict of arrays.