def compute_backward_sinc(dtype, ndim, req):
A = tvm.placeholder([tvm.var() for _ in range(ndim)], name='A', dtype=dtype)
B = tvm.placeholder([tvm.var() for _ in range(ndim)], name='B', dtype=dtype)
C = tvm.placeholder([tvm.var() for _ in range(ndim)], name='C', dtype=dtype)
var = tvm.const(np.pi, dtype)
D = tvm.compute([tvm.var() for _ in range(ndim)],
lambda *index: tvm.if_then_else(B[index] == 0, tvm.const(0, dtype),
(tvm.cos(var * B[index]) / B[index] -
C[index] / B[index]) * A[index]), name='in_grad')
in_grad_a, in_grad = assign_by_req(D, req)
s = tvm.create_schedule(in_grad.op)
s[D].compute_inline()
return s, A, B, C, in_grad_a, in_grad
def cos(x):
"""Take cos of input x.
Parameters
----------
x : tvm.Tensor
Input argument.
Returns
-------
y : tvm.Tensor
The result.
"""
return tvm.compute(x.shape, lambda *i: tvm.cos(x(*i)))