def test_isnan(
low,
high,
shape=(20, 3),
dtype=tvm.float32,
check_round=False,
skip_name_check=False,
):
m = tvm.var("m")
l = tvm.var("l")
A = tvm.placeholder((m, l), dtype=dtype, name="A")
B = topi.isnan(A)
assert tuple(B.shape) == tuple(A.shape)
if not skip_name_check:
assert B.op.body[0].name == "isnan"
a_np = np.random.uniform(low=low, high=high, size=shape).astype(
A.dtype) * 10
a_np.ravel()[np.random.choice(a_np.size,
int(a_np.size * 0.5),
replace=False)] = np.nan
# avoid round check too close to boundary
if check_round:
a_np += ((np.fmod(a_np, 1) - 0.5) < 1e-6) * 1e-5
b_np = np.isnan(a_np)
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 = topi.generic.schedule_injective(B)
foo = tvm.build(s, [A, B], device, name="isnan")
a = tvm.nd.array(a_np, ctx)
b = tvm.nd.array(np.zeros_like(b_np), ctx)
foo(a, b)
tvm.testing.assert_allclose(b.asnumpy(),
b_np,
rtol=1e-5,
atol=1e-5)
check_device('llvm')
check_device('cuda')
check_device('opencl')
check_device('metal')
check_device('rocm')
check_device('vulkan')
check_device('nvptx')
check_device('llvm -device=arm-cpu')
check_device('opencl -device=mali')
check_device('aocl_sw_emu')
def test_isnan(
low,
high,
shape=(20, 3),
dtype="float32",
check_round=False,
skip_name_check=False,
):
m = te.var("m")
l = te.var("l")
A = te.placeholder((m, l), dtype=dtype, name="A")
B = topi.isnan(A)
assert tuple(B.shape) == tuple(A.shape)
if not skip_name_check:
assert B.op.body[0].name == "isnan"
a_np = np.random.uniform(low=low, high=high, size=shape).astype(
A.dtype) * 10
a_np.ravel()[np.random.choice(a_np.size,
int(a_np.size * 0.5),
replace=False)] = np.nan
# avoid round check too close to boundary
if check_round:
a_np += ((np.abs(np.fmod(a_np, 1)) - 0.5) < 1e-6) * 1e-5
b_np = np.isnan(a_np)
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 = topi.testing.get_injective_schedule(device)(B)
foo = tvm.build(s, [A, B], device, name="isnan")
a = tvm.nd.array(a_np, ctx)
b = tvm.nd.array(np.zeros_like(b_np), ctx)
foo(a, b)
tvm.testing.assert_allclose(b.asnumpy(),
b_np,
rtol=1e-5,
atol=1e-5)
for target in get_all_backend():
check_device(target)