def compute_sinc(dtype, ndim):
A = tvm.placeholder([tvm.var() for _ in range(ndim)], name='input', dtype=dtype)
if dtype in ['float16', 'float32', 'float64']:
var = tvm.const(np.pi, dtype)
B = tvm.compute([tvm.var() for _ in range(ndim)],
lambda *index: tvm.if_then_else(A[index] == 0, tvm.const(1, dtype),
tvm.sin(var * A[index]) / (A[index] * var)),
name='output')
else:
var = tvm.const(np.pi, "float64")
B = tvm.compute([tvm.var() for _ in range(ndim)],
lambda *index: tvm.if_then_else(A[index] == 0, tvm.const(1, 'float64'),
tvm.sin(var * A[index].astype('float64')) /
(A[index].astype("float64") * var)),
name='output')
s = tvm.create_schedule(B.op)
return s, A, B
def sin(x):
"""Take sin of input x.
Parameters
----------
x : tvm.Tensor
Input argument.
Returns
-------
y : tvm.Tensor
The result.
"""
return tvm.compute(x.shape, lambda *i: tvm.sin(x(*i)))