Skip to content
1 change: 1 addition & 0 deletions python/tvm/relay/transform/fake_quantization_to_integer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
13 changes: 13 additions & 0 deletions tests/python/relay/test_pass_fake_quantization_to_integer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])