SVD affine is a low rank approximation of the affine layer. It can
    be seen as two consecutive affine layers with a bottleneck. It
    computes:

    .. math::
        {\\mathbf y} = {\\mathbf U} {\\mathbf V} {\\mathbf x} + {\\mathbf b}.

    where :math:`{\\mathbf x}, {\\mathbf y}` are the inputs and
    outputs respectively, and :math:`{\\mathbf U}, {\\mathbf V},
    {\\mathbf b}` are constants.

    The weights :math:`{\\mathbf U}` and :math:`{\\mathbf V}` are
    approximated with singular value decomposition (SVD) of the
    original weight matrix :math:`{\\mathbf W}` and by selecting the
    :math:`{R}` dominant singular values and the corresponding
    singular vectors. Therefore the low rank :math:`{R}` is the size
    of the bottleneck.

    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 weight of the affine is of :math:`{I \\times O}` and
    the compression rate you want to specify is :math:`{CR}`, then you
    set :math:`{R}` as

    .. math::

        R = \\left\\lfloor \\frac{(1 - CR)OI}{O + I} \\right\\rfloor.

    Args:

        inp (~nnabla.Variable): Input N-D array with shape (:math:`M_0
          \\times \ldots \\times M_{B-1} \\times D_B \\times \ldots
          \\times D_N`). Dimensions before and after base_axis are
          flattened as if it is a matrix.

        n_outmaps (int or tuple): Number of output neurons per data.

        r (int): rank of the factorized layer (size of the bottleneck)

        base_axis (int): Dimensions up to `base_axis` are treated as
          the sample 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`.
        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:
        ~nnabla.Variable: :math:`(B + 1)`-D array.
        (:math:`M_0 \\times \ldots \\times M_{B-1} \\times L`)

def svd_affine(inp, n_outmaps, r, base_axis=1, uv_init=None,
               b_init=None, fix_parameters=False, rng=None,
               with_bias=True):
    """SVD affine is a low rank approximation of the affine layer. It can
    be seen as two consecutive affine layers with a bottleneck. It
    computes:

    .. math::
        {\\mathbf y} = {\\mathbf U} {\\mathbf V} {\\mathbf x} + {\\mathbf b}.

    where :math:`{\\mathbf x}, {\\mathbf y}` are the inputs and
    outputs respectively, and :math:`{\\mathbf U}, {\\mathbf V},
    {\\mathbf b}` are constants.

    The weights :math:`{\\mathbf U}` and :math:`{\\mathbf V}` are
    approximated with singular value decomposition (SVD) of the
    original weight matrix :math:`{\\mathbf W}` and by selecting the
    :math:`{R}` dominant singular values and the corresponding
    singular vectors. Therefore the low rank :math:`{R}` is the size
    of the bottleneck.

    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 weight of the affine is of :math:`{I \\times O}` and
    the compression rate you want to specify is :math:`{CR}`, then you
    set :math:`{R}` as

    .. math::

        R = \\left\\lfloor \\frac{(1 - CR)OI}{O + I} \\right\\rfloor.

    Args:

        inp (~nnabla.Variable): Input N-D array with shape (:math:`M_0
          \\times \ldots \\times M_{B-1} \\times D_B \\times \ldots
          \\times D_N`). Dimensions before and after base_axis are
          flattened as if it is a matrix.

        n_outmaps (int or tuple): Number of output neurons per data.

        r (int): rank of the factorized layer (size of the bottleneck)

        base_axis (int): Dimensions up to `base_axis` are treated as
          the sample 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`.
        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:
        ~nnabla.Variable: :math:`(B + 1)`-D array.
        (:math:`M_0 \\times \ldots \\times M_{B-1} \\times L`)

    """
    if not hasattr(n_outmaps, '__iter__'):
        n_outmaps = [n_outmaps]

    n_outmaps = list(n_outmaps)
    n_outmap = int(np.prod(n_outmaps))
    inmaps = np.prod(inp.shape[base_axis:])

    if uv_init is None:
        uv_init = UniformInitializer(
            calc_uniform_lim_glorot(inmaps, n_outmap), 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([int(np.prod(inp.shape[base_axis:])), ] +
                     list(n_outmaps))

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

    if (u is None) or (v is None):
        assert r > 0, "svd_ffine: The rank must larger than zero"
        u_, s_, v_ = np.linalg.svd(uv.reshape(inmaps, n_outmap),
                                   full_matrices=False)
        u_ = np.dot(u_, np.diag(s_))  # fold s into u
        u_ = u_[:, :r]
        v_ = v_[:r, :]
        v_ = v_.reshape([r] + n_outmaps)

        u = nn.Variable([int(np.prod(inp.shape[base_axis:])), r],
                        need_grad=True)
        u.d = u_
        nn.parameter.set_parameter("U", u)

        v = nn.Variable([r] + n_outmaps, need_grad=True)
        v.d = v_
        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)
        v.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", n_outmaps, b_init, True, not fix_parameters)

    return F.affine(F.affine(inp, u, bias=None, base_axis=base_axis),
                    v, bias=b, base_axis=base_axis)