autopdex.utility.mask_select

autopdex.utility.mask_select(array, selection)

Selects elements from a JAX array or dictionary of arrays based on a boolean mask.

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:

Selected elements as an array or dictionary of arrays.