def test_dyn_broadcast_to():
dtype = 'uint8'
rank = 3
shape_type = 'int64'
dyn_shape = relay.Var("shape", relay.ty.TensorType((rank, ), shape_type))
x_shape = (1, )
x = relay.Var("x", relay.ty.TensorType(x_shape, dtype))
z = relay.broadcast_to(x, dyn_shape)
zz = run_infer_type(z)
assert zz.checked_type == relay.ty.TensorType((relay.Any(), ) * rank,
dtype)
func = relay.Function([x, dyn_shape], z)
x = np.random.uniform(size=x_shape).astype(dtype)
dyn_shape = (1, ) * rank
ref_res = np.broadcast_to(x, dyn_shape)
for target, ctx in tvm.testing.enabled_targets():
if (target != 'cuda'
): #skip cuda because we don't have dynamic support for GPU
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind,
mod=mod,
ctx=ctx,
target=target)
op_res = intrp.evaluate(func)(
x, np.array(dyn_shape).astype(shape_type))
tvm.testing.assert_allclose(op_res.asnumpy(),
ref_res,
rtol=1e-5)
def verify_broadcast_to(x_shape, out_shape):
rank = len(out_shape)
dtype = "float32"
shape_type = "int64"
dyn_shape = relay.Var("shape", relay.ty.TensorType((rank, ),
shape_type))
x = relay.Var("x", relay.ty.TensorType(x_shape, dtype))
z = relay.broadcast_to(x, dyn_shape)
zz = run_infer_type(z)
assert zz.checked_type == relay.ty.TensorType((relay.Any(), ) * rank,
dtype)
func = relay.Function([x, dyn_shape], z)
x = np.random.uniform(size=x_shape).astype(dtype)
ref_res = np.broadcast_to(x, out_shape)
for target, ctx in tvm.testing.enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind,
mod=mod,
ctx=ctx,
target=target)
op_res = intrp.evaluate(func)(
x, np.array(out_shape).astype(shape_type))
tvm.testing.assert_allclose(op_res.asnumpy(),
ref_res,
rtol=1e-5)
def test_dyn_broadcast_to():
dtype = "uint8"
rank = 3
shape_type = "int64"
dyn_shape = relay.Var("shape", relay.ty.TensorType((rank, ), shape_type))
x_shape = (1, )
x = relay.Var("x", relay.ty.TensorType(x_shape, dtype))
z = relay.broadcast_to(x, dyn_shape)
zz = run_infer_type(z)
assert zz.checked_type == relay.ty.TensorType((relay.Any(), ) * rank,
dtype)
func = relay.Function([x, dyn_shape], z)
x = np.random.uniform(size=x_shape).astype(dtype)
dyn_shape = (1, ) * rank
ref_res = np.broadcast_to(x, dyn_shape)
for target, dev in tvm.testing.enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
op_res = relay.create_executor(
kind, mod=mod, device=dev, target=target).evaluate(func)(
x, np.array(dyn_shape).astype(shape_type))
tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
def verify_more_dynamic_broadcast_to(x_shape, out_shape):
rank = len(out_shape)
dtype = "float32"
shape_type = "int64"
reshape_shape = relay.Var(
"shape", relay.ty.TensorType((len(x_shape), ), shape_type))
broadcast_shape = relay.Var("shape",
relay.ty.TensorType((rank, ), shape_type))
x = relay.Var("x", relay.ty.TensorType((np.prod(x_shape), ), dtype))
r = relay.reshape(x, reshape_shape)
z = relay.broadcast_to(r, broadcast_shape)
func = relay.Function([x, reshape_shape, broadcast_shape], z)
x = np.random.uniform(size=np.prod(x_shape)).astype(dtype)
ref_res = np.broadcast_to(np.reshape(x, x_shape), out_shape)
for target, ctx in tvm.testing.enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind,
mod=mod,
ctx=ctx,
target=target)
op_res = intrp.evaluate(func)(
x, np.array(x_shape).astype(shape_type),
np.array(out_shape).astype(shape_type))
tvm.testing.assert_allclose(op_res.asnumpy(),
ref_res,
rtol=1e-5)
def test_concretize_broadcast_to_like():
data = relay.var("data", shape=(3, ), dtype="float32")
shape_like = relay.var("shape_like", shape=(3, 3, 3), dtype="float32")
expr = relay.broadcast_to_like(data, shape_like)
expected = run_infer_type(relay.broadcast_to(data, (3, 3, 3)))
actual = run_opt_pass(expr, relay.transform.SimplifyExpr())
assert tvm.ir.structural_equal(actual, expected)
def test_broadcast_to_const_shape_int64():
shape_like = relay.const(np.array([1, 5]), dtype="int64")
x = relay.var("x", shape=(1,), dtype="int64")
z = relay.broadcast_to(x, shape=shape_like)
z = relay.sum(z, axis=0)
f = relay.Function([x], z)
x = np.random.randint(10, size=(1,), dtype="int64")
ref_res = np.broadcast_to(x, (5,))
for target, dev in tvm.testing.enabled_targets():
for kind in ["graph", "debug"]:
op_res = relay.create_executor(kind, device=dev, target=target).evaluate(f)(x)
tvm.testing.assert_allclose(op_res.numpy(), ref_res)
def test_broadcast_to():
shape = (4, 1, 6)
shape_like = (3, 4, 5, 6)
dtype = "float32"
x = relay.Var("x", relay.ty.TensorType(shape, dtype))
z = relay.broadcast_to(x, shape=shape_like)
zz = run_infer_type(z)
assert zz.checked_type == relay.ty.TensorType(shape_like, dtype)
func = relay.Function([x], z)
x = np.random.uniform(size=shape).astype(dtype)
ref_res = np.broadcast_to(x, shape_like)
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)
tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
def test_broadcast_to():
shape = (4, 1, 6)
shape_like = (3, 4, 5, 6)
dtype = "float32"
x = relay.Var("x", relay.ty.TensorType(shape , dtype))
z = relay.broadcast_to(x, shape=shape_like)
zz = relay.ir_pass.infer_type(z)
assert zz.checked_type == relay.ty.TensorType(shape_like, dtype)
func = relay.Function([x], z)
x = np.random.uniform(size=shape).astype(dtype)
ref_res = np.broadcast_to(x, shape_like)
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)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_broadcast_to(shape, broadcast_shape):
x = relay.var("x", relay.TensorType(shape, "float32"))
y = relay.var("y", relay.TensorType(broadcast_shape, "float32"))
z = relay.broadcast_to(x, shape=relay.shape_of(y))
func = run_infer_type(relay.Function([x, y], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("broadcast_to")
assert zz.checked_type == relay.ty.TensorType(broadcast_shape, "float32")
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
y_data = np.random.uniform(low=-1, high=1, size=broadcast_shape).astype("float32")
ref_res = np.broadcast_to(x_data, y_data.shape)
verify_func(func2, [x_data, y_data], ref_res)
def test_broadcast_pool2d_shape_int64(executor_kind):
x_shape = (1, 3, 32, 32)
out_shape = (2, 3, 32, 32)
x = relay.var("data", shape=x_shape, dtype="float32")
broadcast_to = relay.broadcast_to(x,
shape=relay.const([2, 3, 32, 32],
dtype="int64"))
pool2d = relay.nn.max_pool2d(broadcast_to,
pool_size=(3, 3),
padding=(1, 1, 1, 1))
sub = relay.subtract(broadcast_to, pool2d)
f = relay.Function([x], sub)
x = np.ones(x_shape).astype("float32")
ref_res = np.zeros(out_shape).astype("float32")
for target, dev in tvm.testing.enabled_targets():
op_res = relay.create_executor(executor_kind,
device=dev,
target=target).evaluate(f)(x)
tvm.testing.assert_allclose(op_res.numpy(), ref_res)
