def _handle_shape(X, t, name):
    """My version of the shape extracting function. """
    if name == 'obs_mat':
        if len(X.shape) == 4:
            return X[t, ...]
        else:
            return X
