[Relay] add python doc for function in ir_pass (#1877)

include/tvm/relay/pass.h

@@ -29,7 +29,7 @@ Expr InferType(const Environment& env, const Expr& e);
 Expr InferType(const Environment& env, const GlobalVar& var, const Function& f);
 
 /*!
- * \brief Check that types are well formed by applying "kinding rules".
+ * \brief Check that types are well kinded by applying "kinding rules".
  *
  * This pass ensures we do not do things that violate the design of the
  * type system when writing down types.
@@ -39,11 +39,12 @@ Expr InferType(const Environment& env, const GlobalVar& var, const Function& f);
  * We check this by ensuring the `dtype` field of a Tensor always contains
  * a data type such as `int`, `float`, `uint`.
  *
- * \param env The global environment.
  * \param t The type to check.
+ * \param env The global environment.
+ *
  * \return true if the rules are satisified otherwise false
  */
-bool KindCheck(const Environment& env, const Type& t);
+bool KindCheck(const Type& t, const Environment& env);
 
 /*! \brief Compare two expressions for structural equivalence.
  *

python/tvm/relay/ir_pass.py

@@ -15,26 +15,91 @@ def infer_type(env, expr):
     Parameters
     ----------
     env : relay.Environment
-        The global environment.
+      The global environment.
 
     expr : relay.Expr
-        The input expression.
+      The input expression.
 
     Returns
     -------
     checked_expr : relay.Expr
-         The checked expression.
+      The checked expression.
     """
     return _ir_pass.infer_type(env, expr)
 
+def well_formed(e):
+    """Check that each Var is only bound once (well formed).
 
-well_formed = _ir_pass.well_formed
+    Parameters
+    ----------
+    e: relay.Expr
+      The input expression
+
+    Returns
+    -------
+    well_form : bool
+      whether the input expression is well formed
+    """
+    return _ir_pass.well_formed(e)
+
+def check_kind(t, env=None):
+    """Check that the type is well kinded.
+    For example, this mean type cannot has tensor of tensor, or is a tuple type of 2 shapes.
+
+    Parameters
+    ----------
+    t: relay.Type
+      The type to check
+
+    env: relay.Environment, optional
+      The global environment
+
+    Returns
+    -------
+    well_kinded : bool
+      whether the input type is well kinded.
+
+    Examples
+    --------
+    .. code:: python
+
+        assert not check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Shape)]))
+        assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Type)]))
+    """
+    if env is not None:
+        return _ir_pass.check_kind(t, env)
+    else:
+        return _ir_pass.check_kind(t)
 
-check_kind = _ir_pass.check_kind
+def free_vars(e):
+    """Get free variables from expression e.
 
-free_vars = _ir_pass.free_vars
+    Parameters
+    ----------
+    e: relay.Expr
+      The input expression
 
-free_type_vars = _ir_pass.free_type_vars
+    Returns
+    -------
+    free : List[relay.Var]
+      the list of free variables
+    """
+    return _ir_pass.free_vars(e)
+
+def free_type_vars(e):
+    """Get free type variables from expression/type e
+
+    Parameters
+    ----------
+    e: relay.Expr/relay.Type
+      The input expression/type
+
+    Returns
+    -------
+    free : List[relay.TypeParam]
+      the list of free type variables
+    """
+    return _ir_pass.free_type_vars(e)
 
 def dead_code_elimination(e):
     """ Remove expressions which does not effect the program result (dead code).
@@ -59,10 +124,10 @@ def alpha_equal(lhs, rhs):
     ----------
     lhs: relay.Expr
       One of the input Expression.
+
     rhs: relay.Expr
       One of the input Expression.
 
-
     Returns
     -------
     result: bool

src/relay/pass/kind_check.cc

@@ -99,7 +99,7 @@ struct KindChecker : TypeVisitor<> {
   }
 };
 
-bool KindCheck(const Environment& env, const Type &t) {
+bool KindCheck(const Type& t, const Environment& env) {
   KindChecker kc;
   return kc.Check(t);
 }
@@ -107,7 +107,7 @@ bool KindCheck(const Environment& env, const Type &t) {
 TVM_REGISTER_API("relay._ir_pass.check_kind")
     .set_body([](TVMArgs args, TVMRetValue* ret) {
       if (args.size() == 1) {
-        *ret = KindCheck(EnvironmentNode::make({}), args[0]);
+        *ret = KindCheck(args[0], EnvironmentNode::make({}));
       } else {
         *ret = KindCheck(args[0], args[1]);
       }