Apply a function that takes a gradient matrix to a sequence of 2 or 3
    dimensional gradients.

    This is partucularly useful when the gradient of a basis concatenation
    object is quite complex, eg.

    >>> X = np.random.randn(100, 3)
    >>> y = np.random.randn(100)
    >>> N, d = X.shape
    >>> base = RandomRBF(Xdim=d, nbases=5) + RandomRBF(Xdim=d, nbases=5,
    ... lenscale=Parameter(np.ones(d), Positive()))
    >>> Phi = base.transform(X, 1., np.ones(d))
    >>> dffun = lambda dPhi: y.dot(Phi).dot(dPhi.T).dot(y)
    >>> df = apply_grad(dffun, base.grad(X, 1., np.ones(d)))
    >>> np.isscalar(df[0])
    True
    >>> df[1].shape
    (3,)

    Parameters
    ----------
    fun: callable
        the function too apply to the (2d) gradient.
    grad: ndarray or generator
        the gradient of the basis function (output of base.grad).

    Returns
    -------
    scalar, ndarray or sequence:
        the result of applying fun(grad) for a structured grad.

def apply_grad(fun, grad):
    """
    Apply a function that takes a gradient matrix to a sequence of 2 or 3
    dimensional gradients.

    This is partucularly useful when the gradient of a basis concatenation
    object is quite complex, eg.

    >>> X = np.random.randn(100, 3)
    >>> y = np.random.randn(100)
    >>> N, d = X.shape
    >>> base = RandomRBF(Xdim=d, nbases=5) + RandomRBF(Xdim=d, nbases=5,
    ... lenscale=Parameter(np.ones(d), Positive()))
    >>> Phi = base.transform(X, 1., np.ones(d))
    >>> dffun = lambda dPhi: y.dot(Phi).dot(dPhi.T).dot(y)
    >>> df = apply_grad(dffun, base.grad(X, 1., np.ones(d)))
    >>> np.isscalar(df[0])
    True
    >>> df[1].shape
    (3,)

    Parameters
    ----------
    fun: callable
        the function too apply to the (2d) gradient.
    grad: ndarray or generator
        the gradient of the basis function (output of base.grad).

    Returns
    -------
    scalar, ndarray or sequence:
        the result of applying fun(grad) for a structured grad.
    """
    if issequence(grad):
        fgrad = [apply_grad(fun, g) for g in grad]
        return fgrad if len(fgrad) != 1 else fgrad[0]
    elif len(grad) == 0:
        return []
    elif (grad.ndim == 1) or (grad.ndim == 2):
        return fun(grad)
    elif grad.ndim == 3:
        return np.array([fun(grad[:, :, i]) for i in range(grad.shape[2])])
    else:
        raise ValueError("Only up to 3d gradients allowed!")