SVD convolution is a low rank approximation of the convolution
    layer. It can be seen as a depth wise convolution followed by a
    1x1 convolution.

    The flattened kernels for the i-th input map are expressed by
    their low rank approximation. The kernels for the i-th input
    :math:`{\\mathbf W_i}` are approximated with the singular value
    decomposition (SVD) and by selecting the :math:`{R}` dominant
    singular values and the corresponding singular vectors.

    .. math::
        {\\mathbf W_{:,i,:}} ~ {\\mathbf U_i} {\\mathbf V_i}.

    :math:`{\\mathbf U}` contains the weights of the depthwise
    convolution with multiplier :math:`{R}` and :math:`{\\mathbf V}`
    contains the weights of the 1x1 convolution.

    If `uv_init` is a numpy array, :math:`{\\mathbf U}` and
    :math:`{\\mathbf V}` are computed such that `uv_init` is
    approximated by :math:`{\\mathbf{UV}}`. If `uv_init` is `None` or
    an initializer, the product of :math:`{\\mathbf U}` and
    :math:`{\\mathbf V}` approximates the random initialization.

    If :math:`{\\mathbf U}` and :math:`{\\mathbf V}` exist in the
    context, they take precedence over `uv_init`.

    Suppose the kernel tensor of the convolution is of :math:`{O \\times I \\times K \\times K}` and
    the compression rate you want to specify is :math:`{CR}`, then you
    set :math:`{R}` as

    .. math::

        R = \\left\\lfloor \\frac{(1 - CR)OIK^2}{I(O + K^2)} \\right\\rfloor.

    Args:
        inp (~nnabla.Variable): N-D array.

        outmaps (int): Number of convolution kernels (which is equal
          to the number of output channels). For example, to apply
          convolution on an input with 16 types of filters, specify
          16.

        kernel (tuple): Convolution kernel size. For example,
          to apply convolution on an image with a 3 (height) by 5
          (width) two-dimensional kernel, specify (3, 5).

        r (int): Rank of the factorized layer.

        pad (tuple): Padding sizes (`int`) for dimensions.

        stride (tuple): Stride sizes (`int`) for dimensions.

        dilation (tuple): Dilation sizes (`int`) for dimensions.

        uv_init (:obj:`nnabla.initializer.BaseInitializer` or :obj:`numpy.ndarray`):
          Initializer for weight. By default, it is initialized with :obj:`nnabla.initializer.UniformInitializer` within the range determined by :obj:`nnabla.initializer.calc_uniform_lim_glorot`. 

        b_init (:obj:`nnabla.initializer.BaseInitializer` or :obj:`numpy.ndarray`): Initializer for bias. By default, it is initialized with zeros if `with_bias` is `True`.

        base_axis (int): Dimensions up to `base_axis` are treated as the
          sample dimensions.

        fix_parameters (bool): When set to `True`, the weights and
          biases will not be updated.

        rng (numpy.random.RandomState): Random generator for Initializer.

        with_bias (bool): Specify whether to include the bias term.

    Returns:
        :class:`~nnabla.Variable`: :math:`(B + 1)`-D array.
        (:math:`M_0 \\times \ldots \\times M_{B-1} \\times L`)

def svd_convolution(inp, outmaps, kernel, r, pad=None, stride=None,
                    dilation=None, uv_init=None, b_init=None, base_axis=1,
                    fix_parameters=False, rng=None, with_bias=True):
    """SVD convolution is a low rank approximation of the convolution
    layer. It can be seen as a depth wise convolution followed by a
    1x1 convolution.

    The flattened kernels for the i-th input map are expressed by
    their low rank approximation. The kernels for the i-th input
    :math:`{\\mathbf W_i}` are approximated with the singular value
    decomposition (SVD) and by selecting the :math:`{R}` dominant
    singular values and the corresponding singular vectors.

    .. math::
        {\\mathbf W_{:,i,:}} ~ {\\mathbf U_i} {\\mathbf V_i}.

    :math:`{\\mathbf U}` contains the weights of the depthwise
    convolution with multiplier :math:`{R}` and :math:`{\\mathbf V}`
    contains the weights of the 1x1 convolution.

    If `uv_init` is a numpy array, :math:`{\\mathbf U}` and
    :math:`{\\mathbf V}` are computed such that `uv_init` is
    approximated by :math:`{\\mathbf{UV}}`. If `uv_init` is `None` or
    an initializer, the product of :math:`{\\mathbf U}` and
    :math:`{\\mathbf V}` approximates the random initialization.

    If :math:`{\\mathbf U}` and :math:`{\\mathbf V}` exist in the
    context, they take precedence over `uv_init`.

    Suppose the kernel tensor of the convolution is of :math:`{O \\times I \\times K \\times K}` and
    the compression rate you want to specify is :math:`{CR}`, then you
    set :math:`{R}` as

    .. math::

        R = \\left\\lfloor \\frac{(1 - CR)OIK^2}{I(O + K^2)} \\right\\rfloor.

    Args:
        inp (~nnabla.Variable): N-D array.

        outmaps (int): Number of convolution kernels (which is equal
          to the number of output channels). For example, to apply
          convolution on an input with 16 types of filters, specify
          16.

        kernel (tuple): Convolution kernel size. For example,
          to apply convolution on an image with a 3 (height) by 5
          (width) two-dimensional kernel, specify (3, 5).

        r (int): Rank of the factorized layer.

        pad (tuple): Padding sizes (`int`) for dimensions.

        stride (tuple): Stride sizes (`int`) for dimensions.

        dilation (tuple): Dilation sizes (`int`) for dimensions.

        uv_init (:obj:`nnabla.initializer.BaseInitializer` or :obj:`numpy.ndarray`):
          Initializer for weight. By default, it is initialized with :obj:`nnabla.initializer.UniformInitializer` within the range determined by :obj:`nnabla.initializer.calc_uniform_lim_glorot`. 

        b_init (:obj:`nnabla.initializer.BaseInitializer` or :obj:`numpy.ndarray`): Initializer for bias. By default, it is initialized with zeros if `with_bias` is `True`.

        base_axis (int): Dimensions up to `base_axis` are treated as the
          sample dimensions.

        fix_parameters (bool): When set to `True`, the weights and
          biases will not be updated.

        rng (numpy.random.RandomState): Random generator for Initializer.

        with_bias (bool): Specify whether to include the bias term.

    Returns:
        :class:`~nnabla.Variable`: :math:`(B + 1)`-D array.
        (:math:`M_0 \\times \ldots \\times M_{B-1} \\times L`)

    """
    assert r > 0, "svd_convolution: The rank must larger than zero"

    if uv_init is None:
        uv_init = UniformInitializer(
            calc_uniform_lim_glorot(inp.shape[base_axis], outmaps,
                                    tuple(kernel)), rng=rng)

    if type(uv_init) is np.ndarray:
        # TODO: Assert that size of uv_init is correct
        # uv is initialize with numpy array
        uv = uv_init
    else:
        # uv is initialize from initializer
        uv = uv_init((outmaps, inp.shape[base_axis]) + tuple(kernel))

    # flatten kernels
    uv = uv.reshape((outmaps, inp.shape[base_axis], np.prod(kernel)))

    u = get_parameter('U')
    v = get_parameter('V')

    if (u is None) or (v is None):

        inmaps = inp.shape[base_axis]
        u_low_rank = np.zeros((inmaps, np.prod(kernel), r))
        v_low_rank = np.zeros((inmaps, r, outmaps))

        for i in range(inmaps):
            K = np.transpose(uv[:, i, :])
            u_, s_, v_ = np.linalg.svd(K, full_matrices=False)
            u_low_rank[i, :, :] = np.dot(u_[:, :r], np.diag(s_[:r]))
            v_low_rank[i, :, :] = v_[:r, :]

        # reshape U : (I,K*K,r) -> (I*r,K,K) for depthwise conv
        u = nn.Variable((inmaps * r,) + tuple(kernel),
                        need_grad=True)

        u.d = (np.transpose(u_low_rank, axes=(0, 2, 1))
               .reshape((inmaps * r,) + tuple(kernel)))

        nn.parameter.set_parameter("U", u)

        # reshape V :  (I,r,O) -> (O,I*r,1,1) for 1X1 conv
        kernel_one = (1,) * len(kernel)  # 1x1 for 2D convolution
        v = nn.Variable((outmaps, inmaps * r) + kernel_one,
                        need_grad=True)

        v.d = (np.transpose(v_low_rank, axes=(2, 0, 1))
               .reshape((outmaps, inmaps * r) + kernel_one))

        nn.parameter.set_parameter("V", v)

    if fix_parameters == u.need_grad:
        u = u.get_unlinked_variable(need_grad=not fix_parameters)
    if fix_parameters == v.need_grad:
        v = v.get_unlinked_variable(need_grad=not fix_parameters)
    if with_bias and b_init is None:
        b_init = ConstantInitializer()
    b = None
    if with_bias:
        b = get_parameter_or_create(
            "b", (outmaps,), b_init, True, not fix_parameters)

    y = F.depthwise_convolution(inp, u, bias=None, base_axis=base_axis,
                                pad=pad, stride=stride, dilation=dilation,
                                multiplier=r)

    y = F.convolution(y, v, bias=b, base_axis=base_axis, pad=None,
                      stride=None, dilation=None, group=1)
    return y