diff --git a/python/tvm/relay/transform/fake_quantization_to_integer.py b/python/tvm/relay/transform/fake_quantization_to_integer.py index 6032dbf92d..0ed75191c4 100644 --- a/python/tvm/relay/transform/fake_quantization_to_integer.py +++ b/python/tvm/relay/transform/fake_quantization_to_integer.py @@ -87,6 +87,7 @@ def identity(expr, type_map): register_unary_identity("expand_dims") register_unary_identity("nn.max_pool2d") register_unary_identity("nn.batch_flatten") +register_unary_identity("nn.depth_to_space") @register_fake_quantization_to_integer("nn.avg_pool2d") diff --git a/tests/python/relay/test_pass_fake_quantization_to_integer.py b/tests/python/relay/test_pass_fake_quantization_to_integer.py index 7ede17d07d..3680310b4f 100644 --- a/tests/python/relay/test_pass_fake_quantization_to_integer.py +++ b/tests/python/relay/test_pass_fake_quantization_to_integer.py @@ -501,3 +501,16 @@ def test_fake_quantize_pad(): x_np = np.random.randint(-25, 25, size=[1, 383, 128], dtype="int8") compare_fq_to_int(op, [x_np]) + + +def test_fake_quantize_depth_to_space(): + x = relay.var("x", shape=[1, 3, 224, 224], dtype="int8") + + zero = relay.const(0) + x = relay.qnn.op.dequantize(x, relay.const(2.0), zero) + op = relay.op.nn.depth_to_space(x, 4) + op = relay.qnn.op.quantize(op, relay.const(2.0), zero) + + x_np = np.random.randint(-128, 127, size=[1, 3, 224, 224], dtype="int8") + + compare_fq_to_int(op, [x_np])