def test_matmul():
a = sym.Variable("a", shape=(10, 20))
b = sym.Variable("b", shape=(20, 30))
c = sym.matmul(a, b, name="matmul")
g, ldict = correct_layout(c, {"a": "HW", "b": "WC"})
assert (ldict["a"][0] == "HW")
assert (ldict["b"][0] == "WC")
assert (ldict["matmul"][0] == "HC")
# second pass will insert layout transform
_, ldict = correct_layout(g, {"a": "HW16w", "b": "WC16c"})
assert (ldict["a"][0] == "HW16w")
assert (ldict["a_HW"][0] == "HW")
assert (ldict["b"][0] == "WC16c")
assert (ldict["b_WC"][0] == "WC")
assert (ldict["matmul"][0] == "HC")
a = sym.Variable("a", shape=(20, 10))
c = sym.matmul(a, b, name="matmul", transpose_a=True)
g, ldict = correct_layout(c, {"a": "HW", "b": "HC"})
assert (ldict["a"][0] == "HW")
assert (ldict["b"][0] == "HC")
assert (ldict["matmul"][0] == "WC")
b = sym.Variable("b", shape=(30, 20))
c = sym.matmul(a, b, name="matmul", transpose_b=True)
g, ldict = correct_layout(c, {"a": "HW", "b": "CW"})
assert (ldict["a"][0] == "HW")
assert (ldict["b"][0] == "CW")
assert (ldict["matmul"][0] == "HC")
a = sym.Variable("a", shape=(20, 10))
b = sym.Variable("b", shape=(30, 20))
c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
g, ldict = correct_layout(c, {"a": "HW", "b": "CH"})
assert (ldict["a"][0] == "HW")
assert (ldict["b"][0] == "CH")
assert (ldict["matmul"][0] == "WC")
def test_matmul():
a = sym.Variable("a", shape=(10, 20))
b = sym.Variable("b", shape=(20, 30))
c = sym.matmul(a, b, name="matmul")
g, ldict = correct_layout(c, {"a" : "HW", "b" : "WC"})
assert(ldict["a"][0] == "HW")
assert(ldict["b"][0] == "WC")
assert(ldict["matmul"][0] == "HC")
# second pass will insert layout transform
_, ldict = correct_layout(g, {"a" : "HW16w", "b" : "WC16c"})
assert(ldict["a"][0] == "HW16w")
assert(ldict["a_HW"][0] == "HW")
assert(ldict["b"][0] == "WC16c")
assert(ldict["b_WC"][0] == "WC")
assert(ldict["matmul"][0] == "HC")
a = sym.Variable("a", shape=(20, 10))
c = sym.matmul(a, b, name="matmul", transpose_a=True)
g, ldict = correct_layout(c, {"a" : "HW", "b" : "HC"})
assert(ldict["a"][0] == "HW")
assert(ldict["b"][0] == "HC")
assert(ldict["matmul"][0] == "WC")
b = sym.Variable("b", shape=(30, 20))
c = sym.matmul(a, b, name="matmul", transpose_b=True)
g, ldict = correct_layout(c, {"a" : "HW", "b" : "CW"})
assert(ldict["a"][0] == "HW")
assert(ldict["b"][0] == "CW")
assert(ldict["matmul"][0] == "HC")
a = sym.Variable("a", shape=(20, 10))
b = sym.Variable("b", shape=(30, 20))
c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
g, ldict = correct_layout(c, {"a" : "HW", "b" : "CH"})
assert(ldict["a"][0] == "HW")
assert(ldict["b"][0] == "CH")
assert(ldict["matmul"][0] == "WC")
def test_matmul():
a = sym.Variable('a', shape=(10, 20))
b = sym.Variable('b', shape=(20, 30))
c = sym.matmul(a, b, name="matmul")
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
a = sym.Variable('a', shape=(20, 10))
c = sym.matmul(a, b, name="matmul", transpose_a=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
b = sym.Variable('b', shape=(30, 20))
c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
a = sym.Variable('a', shape=(10, 20))
c = sym.matmul(a, b, name="matmul", transpose_b=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
a = sym.Variable('a', shape=(10, 20, 30))
b = sym.Variable('b', shape=(30, 40, 50))
c = sym.matmul(a, b, name="matmul")
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 20, 40, 50])
a = sym.Variable('a', shape=(30, 20, 10))
b = sym.Variable('b', shape=(50, 40, 30))
c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 20, 40, 50])
def test_matmul():
a = sym.Variable('a', shape=(10, 20))
b = sym.Variable('b', shape=(20, 30))
c = sym.matmul(a, b, name="matmul")
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
a = sym.Variable('a', shape=(20, 10))
c = sym.matmul(a, b, name="matmul", transpose_a=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
b = sym.Variable('b', shape=(30, 20))
c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
a = sym.Variable('a', shape=(10, 20))
c = sym.matmul(a, b, name="matmul", transpose_b=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 30])
a = sym.Variable('a', shape=(10, 20, 30))
b = sym.Variable('b', shape=(30, 40, 50))
c = sym.matmul(a, b, name="matmul")
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 20, 40, 50])
a = sym.Variable('a', shape=(30, 20, 10))
b = sym.Variable('b', shape=(50, 40, 30))
c = sym.matmul(a, b, name="matmul", transpose_a=True, transpose_b=True)
sdict = infer_shape(c)
assert(sdict["matmul"][0] == [10, 20, 40, 50])