def METHOD_NAME(labels, match_indices, mask, mask_val):
    """Gather based on batched labels."""
    batch_size = labels.shape[0]
    if batch_size == 1:
        if mask is not None:
            result = _gather_unbatched(
                ops.squeeze(labels, axis=0),
                ops.squeeze(match_indices, axis=0),
                ops.squeeze(mask, axis=0),
                mask_val,
            )
        else:
            result = _gather_unbatched(
                ops.squeeze(labels, axis=0),
                ops.squeeze(match_indices, axis=0),
                None,
                mask_val,
            )
        return ops.expand_dims(result, axis=0)
    else:
        targets = ops.take_along_axis(
            labels, ops.expand_dims(match_indices, axis=-1), axis=1
        )
        if mask is None:
            return targets
        else:
            masked_targets = ops.cast(
                mask_val, labels.dtype
            ) * ops.ones_like(mask, dtype=labels.dtype)
            return ops.where(mask, masked_targets, targets)