utils.batching

Utilities for batched estimators.

select(mask, new, old)

Take the batch elements of new where mask is True and of old otherwise.

Parameters:

Name	Type	Description	Default
mask	Array	Boolean array with the batch shape of the data, e.g., (n_drones,).	required
new	T	Pytree of arrays with leading batch dimensions.	required
old	T	Pytree with the same structure and shapes as new.	required

Source code in drone_estimators/utils/batching.py

def select(mask: Array, new: T, old: T) -> T:
    """Take the batch elements of new where mask is True and of old otherwise.

    Args:
        mask: Boolean array with the batch shape of the data, e.g., (n_drones,).
        new: Pytree of arrays with leading batch dimensions.
        old: Pytree with the same structure and shapes as new.
    """

    def _select(a: Array, b: Array) -> Array:
        xp = array_namespace(mask, a, b)
        return xp.where(xp.reshape(mask, mask.shape + (1,) * (a.ndim - mask.ndim)), a, b)

    return jax.tree.map(_select, new, old)