def test_concretize_collapse_sum_like():
data = relay.var("data", shape=(3, 3, 3), dtype="float32")
shape_like = relay.var("shape_like", shape=(3, ), dtype="float32")
expr = relay.collapse_sum_like(data, shape_like)
expected = run_infer_type(relay.collapse_sum_to(data, (3, )))
actual = run_opt_pass(expr, relay.transform.SimplifyExpr())
assert tvm.ir.structural_equal(actual, expected)
def test_concretize_multiple():
x = relay.var("x", shape=(2, 3), dtype="float32")
y = relay.var("y", shape=(3, ), dtype="float32")
l = x + y
dl = relay.ones_like(l)
dx = relay.zeros_like(x)
dy = relay.zeros_like(y)
dx = dx + relay.collapse_sum_like(dl, dx)
dy = dy + relay.collapse_sum_like(dl, dy)
ret = relay.Tuple([dx, dy])
dl_c = relay.ones((2, 3), "float32")
# NOTE: these are removed by EliminateIdentity
# dx_c = relay.zeros((2, 3), "float32")
# dy_c = relay.zeros((3,), "float32")
dx_c = relay.collapse_sum_to(dl_c, (2, 3))
dy_c = relay.collapse_sum_to(dl_c, (3, ))
ret_c = relay.Tuple([dx_c, dy_c])
expected = run_infer_type(ret_c)
actual = run_opt_pass(ret, relay.transform.SimplifyExpr())
assert tvm.ir.structural_equal(actual, expected)
def test_collapse_sum_like():
shape = (3, 4, 5, 6)
shape_like = (4, 5, 6)
dtype = "float32"
x = relay.Var("x", relay.ty.TensorType(shape, dtype))
y = relay.Var("y", relay.ty.TensorType(shape_like, dtype))
z = relay.collapse_sum_like(x, y)
zz = run_infer_type(z)
assert zz.checked_type == relay.ty.TensorType(shape_like, dtype)
func = relay.Function([x, y], z)
x = np.random.uniform(size=shape).astype(dtype)
y = np.random.uniform(size=shape_like).astype(dtype)
ref_res = np.sum(x, 0)
for target, dev in tvm.testing.enabled_targets():
for kind in ["graph", "debug"]:
op_res = relay.create_executor(kind, device=dev, target=target).evaluate(func)(x, y)
tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
def test_collapse_sum_like():
shape = (3, 4, 5, 6)
shape_like = (4, 5, 6)
dtype = "float32"
x = relay.Var("x", relay.ty.TensorType(shape , dtype))
y = relay.Var("y", relay.ty.TensorType(shape_like, dtype))
z = relay.collapse_sum_like(x, y)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.ty.TensorType(shape_like, dtype)
func = relay.Function([x, y], z)
x = np.random.uniform(size=shape).astype(dtype)
y = np.random.uniform(size=shape_like).astype(dtype)
ref_res = np.sum(x, 0)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x, y)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def test_collapse_sum_like():
x = relay.Var("x", relay.ty.TensorType((3, 4, 5, 6), "int8"))
y = relay.Var("y", relay.ty.TensorType((4, 1, 6), "int8"))
z = relay.collapse_sum_like(x, y)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.ty.TensorType((4, 1, 6), "int8")
