def test_broadcast_to():
x = sym.Variable("x", shape=(4, 1))
y = sym.broadcast_to(x, shape=(0, 4), name="y")
g, ldict = correct_layout(y, "HW")
assert (ldict["x"][0] == "HW")
assert (ldict["y"][0] == "__undef__")
# second pass will insert layout transform
g, ldict = correct_layout(g, "HW16h")
assert (ldict["x"][0] == "HW16h")
assert (ldict["x_HW"][0] == "HW")
assert (ldict["y"][0] == "__undef__")
def test_broadcast_to():
x = sym.Variable("x", shape=(4, 1))
y = sym.broadcast_to(x, shape=(0, 4), name="y")
g, ldict = correct_layout(y, "HW")
assert(ldict["x"][0] == "HW")
assert(ldict["y"][0] == "__undef__")
# second pass will insert layout transform
g, ldict = correct_layout(g, "HW16h")
assert(ldict["x"][0] == "HW16h")
assert(ldict["x_HW"][0] == "HW")
assert(ldict["y"][0] == "__undef__")
def nnvm_distribute(c, v, shp):
"""Implementation of distribute."""
nv = c.ref(v)
assert shp.is_constant()
shp = shp.value
vshp = ashape(v)
if len(shp) != len(vshp):
# We need to pad the shape
nv = sym.expand_dims(nv, axis=0, num_newaxis=len(shp) - len(vshp))
if shp == vshp:
return nv
return sym.broadcast_to(nv, shape=shp)
def check(in_shape, tshape, out_shape):
x = sym.Variable("x", shape=in_shape)
y = sym.broadcast_to(x, shape=tshape, name="y")
sdict = infer_shape(y)
assert(tuple(sdict["y"][0]) == tuple(out_shape))
def test_broadcast_to():
x = sym.Variable('x')
y = sym.broadcast_to(x, shape=(3, 3))
assert y.list_input_names() == ["x"]
def check(in_shape, tshape, out_shape):
x = sym.Variable("x", shape=in_shape)
y = sym.broadcast_to(x, shape=tshape, name="y")
sdict = infer_shape(y)
assert(tuple(sdict["y"][0]) == tuple(out_shape))
def test_broadcast_to():
x = sym.Variable('x')
y = sym.broadcast_to(x, shape=(3, 3))
assert y.list_input_names() == ["x"]
