autopdex.utility.mask_select

autopdex.utility.mask_select(array, selection)[source]

Selects elements from a JAX array or dictionary of arrays based on a boolean mask. Non-selected elements are set to zero.

Parameters:

array (jnp.ndarray or dict) – JAX array or dictionary of JAX arrays.
selection (jnp.ndarray or dict) – Boolean mask or dictionary of boolean masks.

Returns:

Array or dictionary of arrays with non-selected elements set to zero.

Return type:

jnp.ndarray or dict