def test_isnan():
x = tvm.var('x', 'float32')
assert str(tvm.isnan(x)) == 'isnan(x)'
assert str(tvm.isnan(x).dtype) == 'bool'
y = tvm.var('y', 'float16')
assert str(tvm.isnan(y)) == 'isnan(float32(y))'
z = tvm.var('z', 'int32')
assert str(tvm.isnan(z)) == '(bool)0'
k = tvm.var('k', 'int8x2')
assert str(tvm.isnan(k).dtype) == 'uint1x2'
def isnan(x):
"""Check if value of x is NaN, element-wise.
Parameters
----------
x : tvm.Tensor
Input argument.
Returns
-------
y : tvm.Tensor
The result.
"""
return tvm.compute(x.shape, lambda *i: tvm.isnan(x(*i)))