def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
g = dcpe(gradient(f))
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
assert alpha_equal(g, expected)
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
g = dcpe(f, grad=True)
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
expected = run_opt_pass(expected, transform.InferType())
assert tvm.ir.structural_equal(g, expected)
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
g = dcpe(f, grad=True)
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
expected = transform.OptimizeOnExpr(expected, transform.InferType())
assert alpha_equal(g, expected)
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
# TODO(mbs): Revisit once DCE eliminates dead writes.
g = dcpe(f, grad=True, ignore_impurity=True)
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
expected = run_opt_pass(expected, transform.InferType())
assert tvm.ir.structural_equal(g, expected)
def build_impl(self, input_size, memory_size, dtype="float32"):
t = TensorType(shape=(1, memory_size), dtype=dtype)
i = self.input(
var("lstmcell_input", shape=(1, input_size), dtype=dtype))
c = self.input(Var("lstmcell_children", self.p.l(TupleType([t, t]))))
sum = lam(["x", "y"], lambda x, y: x + y)
child_h_sum = self.p.foldl(
sum, op.zeros(shape=(1, memory_size), dtype=dtype),
self.p.map(lam(["z"], lambda z: TupleGetItem(z, 1)), c))
ioux = Linear(input_size=input_size, output_size=memory_size * 3)(i)
iouh = Linear(input_size=memory_size,
output_size=memory_size * 3)(child_h_sum)
iou = ioux + iouh
fx = Linear(input_size=input_size, output_size=memory_size)(i)
fh = Linear(input_size=memory_size, output_size=memory_size)
i, o, u = op.split(iou, 3, axis=1)
i, o, u = op.sigmoid(i), op.sigmoid(o), op.tanh(u)
def foreach_children(children):
f = op.sigmoid(fh(TupleGetItem(children, 1)) + fx)
return f * TupleGetItem(children, 0)
c = self.p.foldl(sum, i * u, self.p.map(lam(["z"], foreach_children),
c))
return Tuple([c, o * op.tanh(c)])
def test_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d * d)
g = dcpe(f, grad=True)
m = d * d
x = relay.Var("x")
o = op.ones_like(x)
x1 = relay.Var("x1")
grad = op.zeros_like(d) + op.collapse_sum_like(x1 * d, d) + op.collapse_sum_like(x1 * d, d)
body = Tuple([x, Tuple([grad])])
body = relay.Let(x1, o, body)
expected = Function([d], relay.Let(x, m, body))
expected = run_opt_pass(expected, transform.InferType())
assert_alpha_equal(g, expected)
def test_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d * d)
# TODO(mbs): Revisit once DCE eliminates dead writes.
g = dcpe(f, grad=True, ignore_impurity=True)
m = d * d
x = relay.Var("x")
o = op.ones_like(x)
x1 = relay.Var("x1")
grad = op.zeros_like(d) + op.collapse_sum_like(x1 * d, d) + op.collapse_sum_like(x1 * d, d)
body = Tuple([x, Tuple([grad])])
body = relay.Let(x1, o, body)
expected = Function([d], relay.Let(x, m, body))
expected = run_opt_pass(expected, transform.InferType())
tvm.ir.assert_structural_equal(g, expected)
def build_impl(self, input_size, memory_size, dtype="float32"):
l = self.input(
Var("l", self.p.l(TensorType(shape=(1, input_size), dtype=dtype))))
def f(c, x):
cell = LSTMCell(input_size=input_size,
memory_size=memory_size,
dtype=dtype)
o = cell(x, self.p.cons(c, self.p.nil()))
return Tuple([o, TupleGetItem(o, 1)])
res = self.p.map_accuml(
lam(["c", "x"], f),
Tuple([
op.zeros(shape=(1, memory_size), dtype=dtype),
op.zeros(shape=(1, memory_size), dtype=dtype)
]), l)
return Tuple(
[TupleGetItem(TupleGetItem(res, 0), 1),
TupleGetItem(res, 1)])
def build_impl(self, input_size, memory_size, dtype="float32"):
l = self.input(
Var("l", self.p.l(TensorType(shape=(1, input_size), dtype=dtype))))
cell = LSTMCell(input_size=input_size,
memory_size=memory_size,
dtype=dtype)
return self.p.foldl(
lam(["c", "x"],
lambda c, x: cell(x, self.p.cons(c, self.p.nil()))),
Tuple([
op.zeros(shape=(1, memory_size), dtype=dtype),
op.zeros(shape=(1, memory_size), dtype=dtype)
]), l)
def build_impl(self, input_size, memory_size, dtype="float32"):
l = self.input(
Var("l", self.p.l(TensorType(shape=(1, input_size), dtype=dtype))))
def LSTM(l):
return LSTMTransformer(input_size=input_size,
memory_size=memory_size,
dtype=dtype)(l)
fwd = LSTM(l)
rev = LSTM(self.p.rev(l))
lhs = op.concatenate(
[TupleGetItem(fwd, 0), TupleGetItem(rev, 0)], axis=1)
t = TensorType(shape=(1, memory_size), dtype=dtype)
x = Var("x", TupleType([t, t])) # cannot infer here
rhs = self.p.map(
Function([x],
op.concatenate([TupleGetItem(x, 0),
TupleGetItem(x, 1)],
axis=1)),
self.p.zip(TupleGetItem(fwd, 1), TupleGetItem(rev, 1)))
return Tuple([lhs, rhs])
def f(c, x):
cell = LSTMCell(input_size=input_size,
memory_size=memory_size,
dtype=dtype)
o = cell(x, self.p.cons(c, self.p.nil()))
return Tuple([o, TupleGetItem(o, 1)])