Add/update documentation (#883)

Summary:
Pull Request resolved: #883

Found some outdated docstrings and Python types while reading codebase.

Reviewed By: chenyang78, aakhundov

Differential Revision: D48135593

fbshipit-source-id: e39eb11aa38b25a8a172b2dadc4d0f7e492a75f5

python/aitemplate/compiler/ops/gemm_universal/gemm_rcr_bias.py

@@ -24,6 +24,7 @@
 
 class gemm_rcr_bias(gemm_rcr):
     """GEMM Specialization: GEMM_RCR(A, B) + Bias
+    A[RowMajor], B[ColMajor], Bias[RowMajor], C[RowMajor]
 
     This operator is equivalent to the following pytorch code:
 

python/aitemplate/compiler/ops/gemm_universal/gemm_rrr_bias.py

@@ -23,6 +23,20 @@
 
 
 class gemm_rrr_bias(gemm_rrr):
+    """GEMM Specialization: GEMM_RRR(A, B) + Bias
+    A[RowMajor], B[RowMajor], Bias[RowMajor], C[RowMajor]
+
+    This operator is equivalent to the following pytorch code:
+
+    .. highlight:: python
+    .. code-block:: python
+        A = torch.randn(M, K).cuda().half()
+        B = torch.randn(K, N).cuda().half()
+        Bias = torch.randn(N).cuda().half()
+
+        y = torch.nn.functional.linear(A, B.t(), bias=Bias)
+    """
+
     def __init__(self):
         super().__init__()
         self._attrs["op"] = "gemm_rrr_bias"

python/aitemplate/compiler/ops/tensor/concatenate.py

@@ -116,7 +116,7 @@ def _infer_shapes(self, inputs: List[Tensor], dim) -> List[IntVar]:
                 output_shape.append(output_dim)
         return output_shape
 
-    def __call__(self, inputs: List[Tensor], dim=0) -> List[Tensor]:
+    def __call__(self, inputs: List[Tensor], dim=0) -> Tensor:
         self._attrs["inputs"] = list(inputs)
         self._attrs["input_accessors"] = [
             TensorAccessor(t) for t in self._attrs["inputs"]

python/aitemplate/compiler/ops/tensor/dynamic_slice.py

@@ -133,9 +133,9 @@ def __call__(
         ----------
         x : Tensor
             Input tensor.
-        start_indices : List[int]
+        start_indices : List[Union[IntVar, IntVarTensor, Optional[int]]]
             Similar to PyTorch and numpy, indices can be negative
-        end_indices : List[int]
+        end_indices : List[Union[IntVar, IntVarTensor, Optional[int]]]
             end_index is not included. Similar to PyTorch and
                 numpy, indices can be negative.
 

python/aitemplate/compiler/ops/tensor/permute0213.py

@@ -61,7 +61,7 @@ def _infer_shapes(self, x: Tensor) -> List[IntVar]:
         x_shape = x._attrs["shape"]
         return [x_shape[0], x_shape[2], x_shape[1], x_shape[3]]
 
-    def __call__(self, x: Tensor) -> List[Tensor]:
+    def __call__(self, x: Tensor) -> Tensor:
         """
         Parameters
         ----------

python/aitemplate/compiler/ops/tensor/permute102.py

@@ -61,7 +61,7 @@ def _infer_shapes(self, x: Tensor) -> List[IntVar]:
         x_shape = x._attrs["shape"]
         return [x_shape[1], x_shape[0], x_shape[2]]
 
-    def __call__(self, x: Tensor) -> List[Tensor]:
+    def __call__(self, x: Tensor) -> Tensor:
         """
         Parameters
         ----------

python/aitemplate/compiler/ops/tensor/permute210.py

@@ -21,7 +21,7 @@
 from aitemplate import backend
 
 from aitemplate.backend import registry
-from aitemplate.compiler.base import Operator, Tensor
+from aitemplate.compiler.base import IntVar, Operator, Tensor
 
 # pylint: disable=C0103,W0221
 
@@ -55,7 +55,7 @@ def __init__(self):
         super().__init__()
         self._attrs["op"] = "permute210"
 
-    def _infer_shapes(self, x: Tensor):
+    def _infer_shapes(self, x: Tensor) -> List[IntVar]:
         """Infers shapes for permute210.
 
         Parameters
@@ -64,15 +64,15 @@ def _infer_shapes(self, x: Tensor):
 
         Returns
         ------
-        Tensor
+        List[IntVar]
             Inferred output 3d tensor with input shape.
             Because its a permute210 operation,
             Out.d0=In.d2, Out.d2=In.d0.
         """
         x_shape = x._attrs["shape"]
         return [x_shape[2], x_shape[1], x_shape[0]]
 
-    def __call__(self, x: Tensor) -> List[Tensor]:
+    def __call__(self, x: Tensor) -> Tensor:
         """
         Return the output tensor of permute210
 

python/aitemplate/compiler/transform/mark_param_tensor.py

@@ -23,6 +23,16 @@
 
 
 def mark_special_views(sorted_graph: List[Tensor]):
+    """
+    Associate each tensor with an external tensor if any of the conditions are true:
+    1. The tensor is a view-of-a-view of an external tensor.
+    2. The tensor is a view of an input, constant or output tensor (i.e. external tensor).
+
+    Parameters
+    ----------
+    sorted_graph : List[Tensor]
+        The graph to mutate.
+    """
     for node in sorted_graph:
         view = node._attrs["is_view_of"]
         if view is None: