from chainer import cuda
from chainer import function
from chainer.utils import type_check


class ExpandDims(function.Function):

    """Expands dimenstions of an input array without copy."""

    def __init__(self, axis):
        self.axis = int(axis)

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 1)
        x_type, = in_types
        if self.axis >= 0:
            type_check.expect(x_type.ndim >= self.axis)
        else:
            type_check.expect(x_type.ndim >= -self.axis - 1)

    def forward(self, x):
        xp = cuda.get_array_module(*x)
        return xp.expand_dims(x[0], self.axis),

    def backward(self, x, gy):
        return gy[0].reshape(x[0].shape),


def expand_dims(x, axis):
    """Expands dimensions of an input variable without copy.

    Args:
        x (~chainer.Variable): Input variable.
        axis (int): Position where new axis is to be inserted.

    Returns:
        ~chainer.Variable: Variable that holds a expanded input.
    """
    return ExpandDims(axis)(x)

@@ -824,16 +824,24 @@
      x (
+:class:%60
 ~chainer
@@ -849,16 +849,78 @@
 Variable
+%60 or :class:%60numpy.ndarray%60 or %5C%0A        :class:%60cupy.ndarray%60
 ): Input
@@ -949,16 +949,28 @@
 s (int):
+%0A           
  Positio
@@ -1008,87 +1008,910 @@
 ted.
-%0A%0A    Returns:%0A        ~chainer.Variable: Variable that holds a expanded input.
+ The %60%60axis%60%60 parameter%0A            is acceptable when :math:%60-ndim - 1 %5C%5Cleq axis %5C%5Cleq ndim%60.%0A            (%60%60ndim%60%60 is the dimension of input variables). When%0A            :math:%60axis %3C 0%60, the result is the same with%0A            :math:%60ndim + 1 - %7Caxis%7C%60.%0A%0A    Returns:%0A        ~chainer.Variable: Variable that holds a expanded input. The %60%60ndim%60%60%0A        of output is one grater than that of %60%60x%60%60.%0A%0A    .. admonition:: Example%0A%0A        %3E%3E%3E x = np.array(%5B1, 2, 3%5D)%0A        %3E%3E%3E x.shape%0A        (3,)%0A        %3E%3E%3E y = F.expand_dims(x, axis=0)%0A        %3E%3E%3E y.shape%0A        (1, 3)%0A        %3E%3E%3E y.data%0A        array(%5B%5B1, 2, 3%5D%5D)%0A        %3E%3E%3E y = F.expand_dims(x, axis=1)%0A        %3E%3E%3E y.shape%0A        (3, 1)%0A        %3E%3E%3E y.data%0A        array(%5B%5B1%5D,%0A               %5B2%5D,%0A               %5B3%5D%5D)%0A        %3E%3E%3E y = F.expand_dims(x, axis=-2)%0A        %3E%3E%3E y.shape%0A        (1, 3)%0A        %3E%3E%3E y.data%0A        array(%5B%5B1, 2, 3%5D%5D)%0A
 %0A   

Improve docs of expand_dims