def test_squeeze():
verify_squeeze((1, 2, 3, 4), 0)
verify_squeeze((1, 2, 1, 4), None)
verify_squeeze((1, 1, 1, 4), (1, 2))
verify_squeeze((1, 1, 1, 1), None)
# a special case to trigger inline let expression
A = te.placeholder((2, ), "float32", "A")
E = topi.squeeze(A)
C = te.compute((1, ), lambda i: E[(2 * A[0] - 1).astype("int32")])
for device in ["cuda", "opencl"]:
ctx = tvm.context(device, 0)
if tvm.testing.device_enabled(device):
with tvm.target.Target(device):
s = tvm.topi.testing.get_injective_schedule(device)(C)
func = tvm.build(s, [A, C])
a = tvm.nd.array(np.array((1, 2)).astype("float32"), ctx=ctx)
c = tvm.nd.empty((1, ), dtype="float32", ctx=ctx)
func(a, c)
assert c.asnumpy()[0] == 2
def test_squeeze():
verify_squeeze((1, 2, 3, 4), 0)
verify_squeeze((1, 2, 1, 4), None)
verify_squeeze((1, 1, 1, 4), (1, 2))
verify_squeeze((1, 1, 1, 1), None)
# a special case to trigger inline let expression
A = te.placeholder((2,), 'float32', 'A')
E = topi.squeeze(A)
C = te.compute((1,), lambda i: E[(2 * A[0] - 1).astype('int32')])
for device in ['cuda', 'opencl']:
ctx = tvm.context(device, 0)
if ctx.exist:
with tvm.target.create(device):
s = tvm.topi.testing.get_injective_schedule(device)(C)
func = tvm.build(s, [A, C])
a = tvm.nd.array(np.array((1, 2)).astype('float32'), ctx=ctx)
c = tvm.nd.empty((1,), dtype='float32', ctx=ctx)
func(a, c)
assert c.asnumpy()[0] == 2
def verify_squeeze(src_shape, axis):
A = te.placeholder(shape=src_shape, name="A")
B = topi.squeeze(A, axis=axis)
def check_device(device, ctx):
print("Running on target: %s" % device)
with tvm.target.Target(device):
s = tvm.topi.testing.get_injective_schedule(device)(B)
foo = tvm.build(s, [A, B], device, name="squeeze")
data_npy = np.random.normal(size=src_shape).astype(A.dtype)
out_npy = np.squeeze(data_npy, axis=axis)
data_nd = tvm.nd.array(data_npy, ctx)
out_nd_shape = out_npy.shape
out_nd = tvm.nd.empty(out_nd_shape, ctx=ctx, dtype=B.dtype)
foo(data_nd, out_nd)
tvm.testing.assert_allclose(out_nd.asnumpy(), out_npy)
for device, ctx in tvm.testing.enabled_targets():
check_device(device, ctx)
def verify_squeeze(src_shape, axis):
A = te.placeholder(shape=src_shape, name="A")
B = topi.squeeze(A, axis=axis)
def check_device(device):
ctx = tvm.context(device, 0)
if not ctx.exist:
print("Skip because %s is not enabled" % device)
return
print("Running on target: %s" % device)
with tvm.target.create(device):
s = tvm.topi.testing.get_injective_schedule(device)(B)
foo = tvm.build(s, [A, B], device, name="squeeze")
data_npy = np.random.normal(size=src_shape).astype(A.dtype)
out_npy = np.squeeze(data_npy, axis=axis)
data_nd = tvm.nd.array(data_npy, ctx)
out_nd_shape = out_npy.shape
out_nd = tvm.nd.empty(out_nd_shape, ctx=ctx, dtype=B.dtype)
foo(data_nd, out_nd)
tvm.testing.assert_allclose(out_nd.asnumpy(), out_npy)
for device in get_all_backend():
check_device(device)