Skip to content
This repository has been archived by the owner on Oct 31, 2023. It is now read-only.

Commit

Permalink
[Relay] fix checkwellform (apache#2705)
Browse files Browse the repository at this point in the history
* do

* address comment
  • Loading branch information
MarisaKirisame authored and bwasti committed Mar 6, 2019
1 parent ea09c28 commit 5471b77
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 5 deletions.
11 changes: 10 additions & 1 deletion src/relay/pass/well_formed.cc
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,15 @@
*/
#include <tvm/relay/pass.h>
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/pattern_functor.h>
#include <unordered_set>

namespace tvm {
namespace relay {


//! brief make sure each Var is bind at most once.
class WellFormedChecker : private ExprVisitor {
class WellFormedChecker : private ExprVisitor, PatternVisitor {
bool well_formed = true;

std::unordered_set<Var, NodeHash, NodeEqual> s;
Expand All @@ -39,6 +40,14 @@ class WellFormedChecker : private ExprVisitor {
CheckWellFormed(f->body);
}

void VisitPattern(const Pattern& p) final {
PatternVisitor::VisitPattern(p);
}

void VisitVar(const Var& v) final {
Check(v);
}

public:
bool CheckWellFormed(const Expr& e) {
this->VisitExpr(e);
Expand Down
27 changes: 23 additions & 4 deletions tests/python/relay/test_ir_well_formed.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import tvm
from tvm import relay
from tvm.relay.ir_pass import well_formed
from tvm.relay.prelude import Prelude

def test_well_formed():
x = relay.Var('x')
def test_let():
x = relay.Var("x")
assert well_formed(x)
v = relay.Constant(tvm.nd.array(10))
ty = None
Expand All @@ -18,7 +19,7 @@ def test_well_formed():


def test_tuple():
x = relay.Var('x')
x = relay.Var("x")
assert well_formed(x)
v = relay.Constant(tvm.nd.array(10))
let = relay.Let(x, v, x)
Expand All @@ -28,5 +29,23 @@ def test_tuple():


def test_tuple_get_item():
t = relay.Var('t')
t = relay.Var("t")
assert well_formed(relay.TupleGetItem(t, 2))


def test_adt():
mod = relay.Module()
p = Prelude(mod)
x = relay.Var("x")
s_case = relay.Clause(relay.PatternConstructor(p.s, [relay.PatternVar(x)]), x)
default_case = relay.Clause(relay.PatternVar(x), x)
m0 = relay.Match(p.z(), [default_case])
m1 = relay.Match(p.z(), [s_case, default_case])
assert well_formed(m0)
assert not well_formed(m1)

if __name__ == "__main__":
test_let()
test_tuple()
test_tuple_get_item()
test_adt()

0 comments on commit 5471b77

Please sign in to comment.