def _compute_intn(value, *indices):
assert input_scale is not None and input_zero_point is not None
# Use indexmod to handle both scalar and per-channel QNN parameters.
scale_idx = tir.indexmod(indices[axis], topi.shape(input_scale)[0])
zp_idx = tir.indexmod(indices[axis], topi.shape(input_zero_point)[0])
return (value[indices] -
input_zero_point[zp_idx]) * input_scale[scale_idx]
def _compute_intn(dtype, value, *indices):
assert output_scale is not None and output_zero_point is not None
const_min = tvm.tir.min_value(dtype)
const_max = tvm.tir.max_value(dtype)
# Use indexmod to handle both scalar and per-channel QNN parameters.
scale_idx = tir.indexmod(indices[axis], topi.shape(output_scale)[0])
zp_idx = tir.indexmod(indices[axis], topi.shape(output_zero_point)[0])
return te.max(
te.min(
te.round(value[indices] / output_scale[scale_idx]) +
output_zero_point[zp_idx],
const_max,
),
const_min,
)
def test_shape():
in_shape = (8, 7, 13)
dtype = "int32"
A = te.placeholder(shape=in_shape, dtype="float32", name="A")
B = topi.shape(A, dtype)
input = np.random.uniform(size=in_shape).astype(A.dtype)
output = np.asarray(in_shape).astype(dtype)
def check_device(device):
ctx = tvm.context(device, 0)
if not ctx.exist:
print("Skip because %s is not enabled" % device)
return
tvm_input = tvm.nd.array(input, ctx)
tvm_output = tvm.nd.empty(output.shape, ctx=ctx, dtype=dtype)
print("Running on target: %s" % device)
with tvm.target.create(device):
s = tvm.topi.testing.get_injective_schedule(device)(B)
f = tvm.build(s, [A, B], device, name="shape")
f(tvm_input, tvm_output)
tvm.testing.assert_allclose(tvm_output.asnumpy(), output)
for backend in get_all_backend():
check_device(backend)
def test_shape():
in_shape = (8, 7, 13)
dtype = "int32"
A = te.placeholder(shape=in_shape, dtype="float32", name="A")
B = topi.shape(A, dtype)
input = np.random.uniform(size=in_shape).astype(A.dtype)
output = np.asarray(in_shape).astype(dtype)
def check_device(device, ctx):
tvm_input = tvm.nd.array(input, ctx)
tvm_output = tvm.nd.empty(output.shape, ctx=ctx, dtype=dtype)
print("Running on target: %s" % device)
with tvm.target.Target(device):
s = tvm.topi.testing.get_injective_schedule(device)(B)
f = tvm.build(s, [A, B], device, name="shape")
f(tvm_input, tvm_output)
tvm.testing.assert_allclose(tvm_output.asnumpy(), output)
for backend, ctx in tvm.testing.enabled_targets():
check_device(backend, ctx)