update python interface

python/paddle/v2/fluid/layers/nn.py

@@ -1702,10 +1702,11 @@ def l2_normalize(x, axis, epsilon=1e-12, name=None):
 
 def matmul(x, y, transpose_x=False, transpose_y=False, name=None):
     """
-    Applies matrix multipication to two tensors. Currently only rank 1 to rank 
-    3 input tensors are supported.
+    Applies matrix multiplication to two tensors. Currently, the input
+    tensors' rank can be any, but when the rank of anyone inputs is
+    bigger than 3, this two inputs' rank should be equal.
 
-    The actual behavior depends on the shapes of :math:`x`, :math:`y` and the 
+    The actual behavior depends on the shapes of :math:`x`, :math:`y` and the
     flag values of :attr:`transpose_x`, :attr:`transpose_y`. Specifically:
 
     - If a transpose flag is specified, the last two dimensions of the tensor 
@@ -1715,17 +1716,17 @@ def matmul(x, y, transpose_x=False, transpose_y=False, name=None):
       opposite: It is treated as :math:`[D, 1]` in nontransposed form and as 
       :math:`[1, D]` in transposed form.
 
-    - After transpose, the two tensors are 2-D or 3-D and matrix multipication 
+    - After transpose, the two tensors are 2-D or n-D and matrix multiplication
       performs in the following way.
 
       - If both are 2-D, they are multiplied like conventional matrices.
-      - If either is 3-D, it is treated as a stack of matrices residing in the 
+      - If either is n-D, it is treated as a stack of matrices residing in the
         last two dimensions and a batched matrix multiply supporting broadcast 
         applies on the two tensors.
 
     Also note that if the raw tensor :math:`x` or :math:`y` is rank-1 and 
     nontransposed, the prepended or appended dimension :math:`1` will be 
-    removed after matrix multipication.
+    removed after matrix multiplication.
 
     Args:
         x (Variable): The input variable which is a Tensor or LoDTensor.
@@ -1742,6 +1743,8 @@ def matmul(x, y, transpose_x=False, transpose_y=False, name=None):
         .. code-block:: python
 
             # Examples to clarify shapes of the inputs and output
+            # x: [B, ..., M, K], y: [B, ..., K, N]
+            fluid.layers.matmul(x, y)  # out: [B, ..., M, N]
             # x: [B, M, K], y: [B, K, N]
             fluid.layers.matmul(x, y)  # out: [B, M, N]
             # x: [B, M, K], y: [K, N]
@@ -1757,9 +1760,9 @@ def matmul(x, y, transpose_x=False, transpose_y=False, name=None):
             fluid.layers.matmul(x, y, True, True)  # out: [M, N]
     """
     helper = LayerHelper('matmul', **locals())
-    assert max(
-        len(x.shape), len(y.shape)
-    ) <= 3, 'Currently only rank 1 to rank 3 input tensors are supported.'
+    assert max(len(x.shape), len(y.shape)) <= 3 or len(x.shape) == len(
+        y.
+        shape), 'Inputs\' rank should be equal or their rank should be less 4.'
     out = helper.create_tmp_variable(dtype=helper.input_dtype())
     helper.append_op(
         type='matmul',