jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_populate_grads_mlp(executor, device, dtype):
    from thunder.benchmarks import NanoGPTMLPBenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility, other settings taken from gpt2 config
    config = NanoGPTConfig(dropout=0, n_layer=12, n_head=12, n_embd=768)

    bench = NanoGPTMLPBenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x,), kwargs = bench.make_batch()

    result = model(x)
    result.backward(torch.ones_like(result))
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)
    tom_grad = grad(tom)
    thunder_grads = tom_grad(x)

    populate_grads(thunder_grads, tom, args=(x,))
    thunder_grads = extract_grads(tom)

    assert_close(torch_grads, thunder_grads, atol=1e-3, rtol=1e-5)
//  [python]
def test_forward_and_backward_from_trace(executor, device, _):
    from thunder import trace
    from thunder.clang import cos, sin
    import thunder.torch as ltorch
    from thunder.core.transforms import forward_and_backward_from_trace, value_and_grad

    def func(a, b, *, c):
        d = a + b + c
        e = d * a + d * b + d * c
        return sin(e) + cos(e), e, ltorch.sin(e) + ltorch.cos(e)

    expected_vjp_func = executor.make_callable_legacy(value_and_grad(func))

    a = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)
    b = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)
    c = make_tensor((3,), device=device, dtype=torch.float64, requires_grad=True)
    trace = trace(inline_trace=False)(func, a, b, c=c)
    fw_trace, bw_trace = forward_and_backward_from_trace(trace)
    fw = executor.make_callable(fw_trace)
    bw = executor.make_callable(bw_trace)
    fw_out, saved_for_backward = fw(a, b, c=c)
    expected_fw_out, expected_grads = expected_vjp_func(a, b, c=c)
    torch.testing.assert_allclose(fw_out, expected_fw_out)

    output_grads = tree_map(lambda x: torch.ones_like(x), fw_out)
    bw_out = bw(saved_for_backward, output_grads)
    torch.testing.assert_allclose(bw_out, expected_grads)
//  [python]
def test_forward_and_backward_from_trace(executor, device, _):
    from thunder import trace
    from thunder.clang import cos, sin
    import thunder.torch as ltorch
    from thunder.core.transforms import forward_and_backward_from_trace, value_and_grad

    def func(a, b, *, c):
        d = a + b + c
        e = d * a + d * b + d * c
        return sin(e) + cos(e), e, ltorch.sin(e) + ltorch.cos(e)

    expected_vjp_func = executor.make_callable_legacy(value_and_grad(func))

    a = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)
    b = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)
    c = make_tensor((3,), device=device, dtype=torch.float64, requires_grad=True)
    trace = trace(inline_trace=False)(func, a, b, c=c)
    fw_trace, bw_trace = forward_and_backward_from_trace(trace)
    fw = executor.make_callable(fw_trace)
    bw = executor.make_callable(bw_trace)
    fw_out, saved_for_backward = fw(a, b, c=c)
    expected_fw_out, expected_grads = expected_vjp_func(a, b, c=c)
    torch.testing.assert_close(fw_out, expected_fw_out)

    output_grads = tree_map(lambda x: torch.ones_like(x), fw_out)
    bw_out = bw(saved_for_backward, output_grads)
    torch.testing.assert_close(bw_out, expected_grads)
//  [python]
def test_populate_grads_csa(executor, device, dtype):
    from thunder.benchmarks import NanoGPTCSABenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility, other settings taken from gpt2 config
    config = NanoGPTConfig(dropout=0, n_layer=12, n_head=12, n_embd=768)

    bench = NanoGPTCSABenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x,), kwargs = bench.make_batch()

    result = model(x)
    result.backward(torch.ones_like(result))
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)
    tom_grad = grad(tom)
    thunder_grads = tom_grad(x)

    populate_grads(thunder_grads, tom, args=[x])
    thunder_grads = extract_grads(tom)

    t.assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
//  [python]
def test_populate_grads_csa(executor, device, dtype):
    from thunder.benchmarks import NanoGPTCSABenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility, other settings taken from gpt2 config
    config = NanoGPTConfig(dropout=0, n_layer=12, n_head=12, n_embd=768)

    bench = NanoGPTCSABenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x,), kwargs = bench.make_batch()

    result = model(x)
    result.backward(torch.ones_like(result))
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)
    tom_grad = grad(tom)
    thunder_grads = tom_grad(x)

    populate_grads(thunder_grads, tom, args=[x])
    thunder_grads = extract_grads(tom)

    assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
//  [python]
def test_populate_grads_block(executor, device, dtype):
    from thunder.benchmarks import NanoGPTBlockBenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility, other settings taken from gpt2 config
    config = NanoGPTConfig(dropout=0, n_layer=12, n_head=12, n_embd=768)

    bench = NanoGPTBlockBenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x,), kwargs = bench.make_batch()

    result = model(x)
    result.backward(torch.ones_like(result))
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)
    tom_grad = grad(tom)
    thunder_grads = tom_grad(x)

    populate_grads(thunder_grads, tom, args=[x])
    thunder_grads = extract_grads(tom)

    t.assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
//  [python]
def test_populate_grads_block(executor, device, dtype):
    from thunder.benchmarks import NanoGPTBlockBenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility, other settings taken from gpt2 config
    config = NanoGPTConfig(dropout=0, n_layer=12, n_head=12, n_embd=768)

    bench = NanoGPTBlockBenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x,), kwargs = bench.make_batch()

    result = model(x)
    result.backward(torch.ones_like(result))
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)
    tom_grad = grad(tom)
    thunder_grads = tom_grad(x)

    populate_grads(thunder_grads, tom, args=[x])
    thunder_grads = extract_grads(tom)

    assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
//  [python]
def test_fori_loop_no_export():
    if not pytorch_pfn_extras.requires("1.8.0"):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    torch.manual_seed(0)

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.relu = nn.ReLU()

        def forward(self, x):
            x = x * 0.5
            h = lax.fori_loop(1, 4, lambda it, val: it * self.relu(self.linear(val)), x)
            return h + 1

    model = Net()
    x = torch.ones((1, 1))
    y = model(x)
    v = x * 0.5
    for i in range(1, 4):
        v = i * model.relu(model.linear(v))
    y_expected = v + 1
    torch.testing.assert_allclose(y, y_expected)
//  [python]
def test_fori_loop_no_export():
    if not pytorch_pfn_extras.requires("1.8.0"):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    torch.manual_seed(0)

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.relu = nn.ReLU()

        def forward(self, x):
            x = x * 0.5
            h = lax.fori_loop(1, 4, lambda it, val: it * self.relu(self.linear(val)), x)
            return h + 1

    model = Net()
    x = torch.ones((1, 1))
    y = model(x)
    v = x * 0.5
    for i in range(1, 4):
        v = i * model.relu(model.linear(v))
    y_expected = v + 1
    torch.testing.assert_close(y, y_expected)
//  [python]
def test_pyl_torsion_angle():
    p = protein_to_pyg(pdb_code="1nth")

    pyl_index = p.residue_id.index("A:PYL:202:")
    sc_angles = sidechain_torsion(p.coords, p.residues)

    t.assert_close(
        torch.zeros_like(sc_angles[pyl_index]),
        sc_angles[pyl_index],
        rtol=1e-5,
        atol=1e-5,
    )
//  [python]
def test_pyl_torsion_angle():
    p = protein_to_pyg(pdb_code="1nth")

    pyl_index = p.residue_id.index("A:PYL:202:")
    sc_angles = sidechain_torsion(p.coords, p.residues)

    torch.testing.assert_close(
        torch.zeros_like(sc_angles[pyl_index]),
        sc_angles[pyl_index],
        rtol=1e-5,
        atol=1e-5,
    )
//  [python]
def test_populate_grads_nanogpt(executor, device, dtype):
    import sys

    if sys.platform == "win32":
        pytest.skip(
            "This test crashes its worked on Windows when run using pytest distributed (Windows fatal exception: access violation)"
        )
    if IN_CI and torch.device(device).type == "cpu":
        pytest.skip("Skipping the CPU version of this test in CI because it's very slow")

    from thunder.benchmarks import NanoGPTBenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility
    config = NanoGPTConfig(dropout=0, n_layer=2, n_head=1, n_embd=64)

    bench = NanoGPTBenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x, targets), kwargs = bench.make_batch()

    logits, loss = model(x, targets)
    loss.backward()
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)

    def grad_specifier(out) -> None:
        logits, loss = out
        put_grad(loss, ltorch.ones_like(loss))

    tom_grad = grad(tom, grad_specifier=grad_specifier)
    thunder_grads = tom_grad(x, targets)

    populate_grads(thunder_grads, tom, args=[x, targets])
    thunder_grads = extract_grads(tom)

    assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
//  [python]
def test_fori_loop():
    if not pytorch_pfn_extras.requires('1.8.0'):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
        pytest.skip('ONNX grad test does not work in windows CI for torch >= 1.10')

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.linear.weight.data[:] = 2
            self.linear.bias.data[:] = 3

        def forward(self, x):
            h = lax.fori_loop(1, 4, lambda it, val: it * self.linear(val) + 0.5, x)
            return h

    model = Net()
    x = torch.tensor([[0], [1]]).float()
    output_dir = _helper(
        model,
        x,
        'fori_loop',
        input_names=("x",),
        enable_onnx_checker=False,
        use_pfto=False,
        do_constant_folding=False,
    )

    actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
    assert len([x for x in actual_onnx.graph.node if x.op_type == "Loop"]) == 1
    ort_session = _ort_session(os.path.join(output_dir, "model.onnx"))
    actual = ort_session.run(None, {"x": x.cpu().numpy()})
    expected = model(x)
    torch.testing.assert_allclose(expected, torch.tensor(actual[0]))
//  [python]
def test_fori_loop():
    if not pytorch_pfn_extras.requires('1.8.0'):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
        pytest.skip('ONNX grad test does not work in windows CI for torch >= 1.10')

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.linear.weight.data[:] = 2
            self.linear.bias.data[:] = 3

        def forward(self, x):
            h = lax.fori_loop(1, 4, lambda it, val: it * self.linear(val) + 0.5, x)
            return h

    model = Net()
    x = torch.tensor([[0], [1]]).float()
    output_dir = _helper(
        model,
        x,
        'fori_loop',
        input_names=("x",),
        enable_onnx_checker=False,
        use_pfto=False,
        do_constant_folding=False,
    )

    actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
    assert len([x for x in actual_onnx.graph.node if x.op_type == "Loop"]) == 1
    ort_session = _ort_session(os.path.join(output_dir, "model.onnx"))
    actual = ort_session.run(None, {"x": x.cpu().numpy()})
    expected = model(x)
    torch.testing.assert_close(expected, torch.tensor(actual[0]))
//  [python]
def torch_assertclose(
    a: torch.Tensor, b: torch.Tensor, *, atol: float = 1e-05, rtol: float = 1e-05
):
    t.assert_close(
        a,
        b,
        atol=atol,
        rtol=rtol,
    )
//  [python]
def torch_assertclose(
    a: torch.Tensor, b: torch.Tensor, *, atol: float = 1e-05, rtol: float = 1e-05
):
    torch.testing.assert_close(
        a,
        b,
        atol=atol,
        rtol=rtol,
    )
//  [python]
def test_while_loop():
    if not pytorch_pfn_extras.requires('1.8.0'):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
        pytest.skip('ONNX grad test does not work in windows CI for torch >= 1.10')

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.linear.weight.data[:] = 2
            self.linear.bias.data[:] = 3

        def forward(self, x):
            def cond_fn(x):
                return x.sum() < 100

            def body_fn(x):
                return self.linear(x)
            h = lax.while_loop(cond_fn, body_fn, x)
            return h

    model = Net()
    x = torch.tensor([[0], [1]]).float()
    output_dir = _helper(
        model,
        x,
        'while_loop',
        input_names=("x",),
        enable_onnx_checker=False,
        use_pfto=False,
        do_constant_folding=False,
    )

    actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
    assert len([x for x in actual_onnx.graph.node if x.op_type == "Loop"]) == 1
    ort_session = _ort_session(os.path.join(output_dir, "model.onnx"))
    actual = ort_session.run(None, {"x": x.cpu().numpy()})
    expected = model(x)
    t.assert_close(expected, torch.tensor(actual[0]))
//  [python]
def test_while_loop():
    if not pytorch_pfn_extras.requires('1.8.0'):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
        pytest.skip('ONNX grad test does not work in windows CI for torch >= 1.10')

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.linear.weight.data[:] = 2
            self.linear.bias.data[:] = 3

        def forward(self, x):
            def cond_fn(x):
                return x.sum() < 100

            def body_fn(x):
                return self.linear(x)
            h = lax.while_loop(cond_fn, body_fn, x)
            return h

    model = Net()
    x = torch.tensor([[0], [1]]).float()
    output_dir = _helper(
        model,
        x,
        'while_loop',
        input_names=("x",),
        enable_onnx_checker=False,
        use_pfto=False,
        do_constant_folding=False,
    )

    actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
    assert len([x for x in actual_onnx.graph.node if x.op_type == "Loop"]) == 1
    ort_session = _ort_session(os.path.join(output_dir, "model.onnx"))
    actual = ort_session.run(None, {"x": x.cpu().numpy()})
    expected = model(x)
    torch.testing.assert_close(expected, torch.tensor(actual[0]))
//  [python]
def test_cond():
    if not pytorch_pfn_extras.requires('1.8.0'):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
        pytest.skip('ONNX grad test does not work in windows CI for torch >= 1.10')

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.linear.weight.data[:] = 2
            self.linear.bias.data[:] = 3

        def forward(self, x):
            def true_fn(x):
                return self.linear(x)

            def false_fn(x):
                return -x

            x = self.linear(x)
            h = lax.cond(x.sum().long() % 2 == 0, true_fn, false_fn, x)
            return h

    model = Net()
    x = torch.tensor([[0], [1]]).float()
    output_dir = _helper(
        model,
        x,
        'cond',
        input_names=("x",),
        enable_onnx_checker=False,
        use_pfto=False,
        do_constant_folding=False,
    )

    actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
    assert len([x for x in actual_onnx.graph.node if x.op_type == "If"]) == 1
    ort_session = _ort_session(os.path.join(output_dir, "model.onnx"))
    actual = ort_session.run(None, {"x": x.cpu().numpy()})
    expected = model(x)
    t.assert_close(expected, torch.tensor(actual[0]))
//  [python]
def test_cond():
    if not pytorch_pfn_extras.requires('1.8.0'):
        pytest.skip('skip for PyTorch 1.7 or earlier')

    if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
        pytest.skip('ONNX grad test does not work in windows CI for torch >= 1.10')

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.linear = nn.Linear(1, 1)
            self.linear.weight.data[:] = 2
            self.linear.bias.data[:] = 3

        def forward(self, x):
            def true_fn(x):
                return self.linear(x)

            def false_fn(x):
                return -x

            x = self.linear(x)
            h = lax.cond(x.sum().long() % 2 == 0, true_fn, false_fn, x)
            return h

    model = Net()
    x = torch.tensor([[0], [1]]).float()
    output_dir = _helper(
        model,
        x,
        'cond',
        input_names=("x",),
        enable_onnx_checker=False,
        use_pfto=False,
        do_constant_folding=False,
    )

    actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
    assert len([x for x in actual_onnx.graph.node if x.op_type == "If"]) == 1
    ort_session = _ort_session(os.path.join(output_dir, "model.onnx"))
    actual = ort_session.run(None, {"x": x.cpu().numpy()})
    expected = model(x)
    torch.testing.assert_close(expected, torch.tensor(actual[0]))
//  [python]
def test_torch_layers_scripting():
    if not is_backend_tested("torch"):
        pytest.skip()
    else:
        import torch

        for script_layer in [False, True]:
            model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer)
            model2 = torch.jit.script(model1)
            input = torch.randn([10, 3, 32, 32])

            torch.testing.assert_allclose(model1(input), model2(input), atol=1e-3, rtol=1e-3)
//  [python]
def test_torch_layers_scripting():
    if not is_backend_tested("torch"):
        pytest.skip()
    else:
        import torch

        for script_layer in [False, True]:
            model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer)
            model2 = torch.jit.script(model1)
            input = torch.randn([10, 3, 32, 32])

            torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3)
//  [python]
def test_fused_obs_fake_quant_backward_op(self, device) -> None:
        n = m = k = 10
        input_shape = (m, n)
        output_shape = (m, n)

        x = torch.randn(input_shape, device=device, requires_grad=True)

        avg_const = 0.01
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        x_min, x_max = _get_tensor_min_max(x)
        x_scale, x_zero_point = _get_scale_zp(
            x_min, x_max, torch.quint8
        )

        x_scale = torch.tensor(x_scale, device=device)
        x_zero_point = torch.tensor(x_zero_point, dtype=torch.int, device=device)
        x_fake_quant = torch.fake_quantize_per_tensor_affine(
            x, x_scale, x_zero_point, 0, 255
        )

        pt_op = torch.fused_moving_avg_obs_fake_quant
        out = pt_op(
            x,
            torch.tensor(1, device=device),
            torch.tensor(1, device=device),
            torch.tensor(x_min, device=device),
            torch.tensor(x_max, device=device),
            scale,
            zero_point,
            avg_const,
            0,
            255,
            0,
            False,
        )
        # verify the output matches
        torch.testing.assert_allclose(out, x_fake_quant)

        # verify the gradient matches expectation of fake_quant op
        dout = torch.rand_like(x, dtype=torch.float).to(device)
        out.backward(dout)

        dX = _fake_quantize_per_tensor_affine_grad_reference(
            dout, x, x_scale, x_zero_point, 0, 255)
        self.assertEqual(dX, x.grad)
        self.assertTrue(x.grad.dtype == torch.float32)
//  [python]
def test_fused_obs_fake_quant_backward_op(self, device) -> None:
        n = m = k = 10
        input_shape = (m, n)
        output_shape = (m, n)

        x = torch.randn(input_shape, device=device, requires_grad=True)

        avg_const = 0.01
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        x_min, x_max = _get_tensor_min_max(x)
        x_scale, x_zero_point = _get_scale_zp(
            x_min, x_max, torch.quint8
        )

        x_scale = torch.tensor(x_scale, device=device)
        x_zero_point = torch.tensor(x_zero_point, dtype=torch.int, device=device)
        x_fake_quant = torch.fake_quantize_per_tensor_affine(
            x, x_scale, x_zero_point, 0, 255
        )

        pt_op = torch.fused_moving_avg_obs_fake_quant
        out = pt_op(
            x,
            torch.tensor(1, device=device),
            torch.tensor(1, device=device),
            torch.tensor(x_min, device=device),
            torch.tensor(x_max, device=device),
            scale,
            zero_point,
            avg_const,
            0,
            255,
            0,
            False,
        )
        # verify the output matches
        torch.testing.assert_close(out, x_fake_quant)

        # verify the gradient matches expectation of fake_quant op
        dout = torch.rand_like(x, dtype=torch.float).to(device)
        out.backward(dout)

        dX = _fake_quantize_per_tensor_affine_grad_reference(
            dout, x, x_scale, x_zero_point, 0, 255)
        self.assertEqual(dX, x.grad)
        self.assertTrue(x.grad.dtype == torch.float32)
//  [python]
def test_fused_backward_op_fake_quant_off(self, device) -> None:
        n = m = 4
        input_shape = (m, n)
        output_shape = (m, n)

        x = torch.randn(input_shape, device=device, requires_grad=True)

        avg_const = 0.01
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        x_min, x_max = _get_tensor_min_max(x)
        x_scale, x_zero_point = _get_scale_zp(
            x_min, x_max, torch.quint8
        )


        pt_op = torch.fused_moving_avg_obs_fake_quant
        out = pt_op(
            x,
            torch.tensor(0, device=device),
            torch.tensor(0, device=device),
            torch.tensor(x_min, device=device),
            torch.tensor(x_max, device=device),
            scale,
            zero_point,
            avg_const,
            0,
            255,
            0,
            False,
        )
        # verify the output matches
        torch.allclose(out, x)

        # verify the gradient matches expectation of fake_quant op
        dout = torch.rand_like(x, dtype=torch.float).to(device)
        out.backward(dout)

        dX = _fake_quantize_per_tensor_affine_grad_reference(
            dout, x, x_scale, x_zero_point, 0, 255)
        self.assertEqual(dX, x.grad)
        self.assertTrue(x.grad.dtype == torch.float32)
//  [python]
def test_fused_backward_op_fake_quant_off(self, device) -> None:
        n = m = 4
        input_shape = (m, n)
        output_shape = (m, n)

        x = torch.randn(input_shape, device=device, requires_grad=True)

        avg_const = 0.01
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        x_min, x_max = _get_tensor_min_max(x)
        x_scale, x_zero_point = _get_scale_zp(
            x_min, x_max, torch.quint8
        )


        pt_op = torch.fused_moving_avg_obs_fake_quant
        out = pt_op(
            x,
            torch.tensor(0, device=device),
            torch.tensor(0, device=device),
            torch.tensor(x_min, device=device),
            torch.tensor(x_max, device=device),
            scale,
            zero_point,
            avg_const,
            0,
            255,
            0,
            False,
        )
        # verify the output matches
        torch.testing.assert_close(out, x)

        # verify the gradient matches expectation of fake_quant op
        dout = torch.rand_like(x, dtype=torch.float).to(device)
        out.backward(dout)

        dX = _fake_quantize_per_tensor_affine_grad_reference(
            dout, x, x_scale, x_zero_point, 0, 255)
        self.assertEqual(dX, x.grad)
        self.assertTrue(x.grad.dtype == torch.float32)
//  [python]
def test_torch_layer():
    if not is_backend_tested("torch"):
        pytest.skip()
    else:
        # checked that torch present
        import torch
        import torch.jit

        model1 = create_torch_model(use_reduce=True)
        model2 = create_torch_model(use_reduce=False)
        input = torch.randn([10, 3, 32, 32])
        # random models have different predictions
        assert not torch.allclose(model1(input), model2(input))
        model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict())))
        assert torch.allclose(model1(input), model2(input))

        # tracing (freezing)
        model3 = torch.jit.trace(model2, example_inputs=input)
        torch.allclose(model1(input), model3(input), atol=1e-3, rtol=1e-3)
        torch.allclose(model1(input + 1), model3(input + 1), atol=1e-3, rtol=1e-3)

        model4 = torch.jit.trace(model2, example_inputs=input)
        torch.allclose(model1(input), model4(input), atol=1e-3, rtol=1e-3)
        torch.allclose(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3)
//  [python]
def test_torch_layer():
    if not is_backend_tested("torch"):
        pytest.skip()
    else:
        # checked that torch present
        import torch
        import torch.jit

        model1 = create_torch_model(use_reduce=True)
        model2 = create_torch_model(use_reduce=False)
        input = torch.randn([10, 3, 32, 32])
        # random models have different predictions
        assert not torch.allclose(model1(input), model2(input))
        model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict())))
        assert torch.allclose(model1(input), model2(input))

        # tracing (freezing)
        model3 = torch.jit.trace(model2, example_inputs=input)
        torch.testing.assert_close(model1(input), model3(input), atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(model1(input + 1), model3(input + 1), atol=1e-3, rtol=1e-3)

        model4 = torch.jit.trace(model2, example_inputs=input)
        torch.testing.assert_close(model1(input), model4(input), atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3)
//  [python]
def run_worker(args):
    mod = ExampleCode()
    mod.to(args.device)

    x = torch.randn(batch_size, d_hid, device=args.device)

    split_policy = split_into_equal_size(args.world_size)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
        split_policy=split_policy,
    )

    # Check returned number of stages
    assert (
        pipe.num_stages == args.world_size
    ), f"Model is split into {pipe.num_stages} stages instead of {args.world_size}"
    print(f"Split test passed: got {pipe.num_stages} stages")

    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Attach to a schedule
    schedule = ScheduleGPipe(stage, args.chunks)

    # Run
    if args.rank == 0:
        schedule.step(x)
    else:
        out = schedule.step()

    dist.barrier()
    print(f"Rank {args.rank} completes")

    # Last rank checks result
    if args.rank == args.world_size - 1:
        ref_out = mod(x)
        torch.testing.assert_allclose(out, ref_out)
        print(
            f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref_out)}"
        )
//  [python]
def run_worker(args):
    mod = ExampleCode()
    mod.to(args.device)

    x = torch.randn(batch_size, d_hid, device=args.device)

    split_policy = split_into_equal_size(args.world_size)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
        split_policy=split_policy,
    )

    # Check returned number of stages
    assert (
        pipe.num_stages == args.world_size
    ), f"Model is split into {pipe.num_stages} stages instead of {args.world_size}"
    print(f"Split test passed: got {pipe.num_stages} stages")

    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Attach to a schedule
    schedule = ScheduleGPipe(stage, args.chunks)

    # Run
    if args.rank == 0:
        schedule.step(x)
    else:
        out = schedule.step()

    dist.barrier()
    print(f"Rank {args.rank} completes")

    # Last rank checks result
    if args.rank == args.world_size - 1:
        ref_out = mod(x)
        torch.testing.assert_close(out, ref_out)
        print(
            f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref_out)}"
        )
//  [python]
def test_conv_bias():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation()
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_allclose(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_allclose(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_allclose(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation()
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_close(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_close(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_add():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_add(
                x, conv.weight, conv.bias, y, alpha, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation()
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    y = torch.ones(1, 1, 254, 254, requires_grad=True).cuda()
    out = aot_printer(model)(x, y, 0.5)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_y = torch.ones(1, 1, 254, 254, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x, fused_y, 0.5)
    fused_out.sum().backward()

    t.assert_allclose(fused_out, out, rtol=1e-3, atol=1e-3)
    t.assert_allclose(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    t.assert_allclose(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_add():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_add(
                x, conv.weight, conv.bias, y, alpha, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation()
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    y = torch.ones(1, 1, 254, 254, requires_grad=True).cuda()
    out = aot_printer(model)(x, y, 0.5)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_y = torch.ones(1, 1, 254, 254, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x, fused_y, 0.5)
    fused_out.sum().backward()

    torch.testing.assert_close(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_close(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_sigmoid():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_sigmoid(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation(activation_cls=torch.nn.Sigmoid)
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    t.assert_close(fused_out, out, rtol=1e-3, atol=1e-3)
    t.assert_close(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    t.assert_close(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_sigmoid():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_sigmoid(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation(activation_cls=torch.nn.Sigmoid)
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_close(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_close(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_relu():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_relu(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_allclose(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_allclose(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_allclose(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_relu():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_relu(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_close(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_close(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_decode_with_exif_orientation(tmpdir, codec, orientation):
    fp = os.path.join(tmpdir, f"exif_oriented_{orientation}.{codec}")
    t = torch.randint(0, 256, size=(3, 256, 257), dtype=torch.uint8)
    im = F.to_pil_image(t)
    exif = im.getexif()
    exif[0x0112] = orientation  # set exif orientation
    im.save(fp, codec.upper(), exif=exif.tobytes())

    data = read_file(fp)
    output = decode_image(data, apply_exif_orientation=True)

    pimg = Image.open(fp)
    pimg = ImageOps.exif_transpose(pimg)

    expected = F.pil_to_tensor(pimg)
    torch.testing.assert_allclose(expected, output)
//  [python]
def test_decode_with_exif_orientation(tmpdir, codec, orientation):
    fp = os.path.join(tmpdir, f"exif_oriented_{orientation}.{codec}")
    t = torch.randint(0, 256, size=(3, 256, 257), dtype=torch.uint8)
    im = F.to_pil_image(t)
    exif = im.getexif()
    exif[0x0112] = orientation  # set exif orientation
    im.save(fp, codec.upper(), exif=exif.tobytes())

    data = read_file(fp)
    output = decode_image(data, apply_exif_orientation=True)

    pimg = Image.open(fp)
    pimg = ImageOps.exif_transpose(pimg)

    expected = F.pil_to_tensor(pimg)
    torch.testing.assert_close(expected, output)
//  [python]
def test_invalid_exif(tmpdir, size):
    # Inspired from a PIL test:
    # https://github.com/python-pillow/Pillow/blob/8f63748e50378424628155994efd7e0739a4d1d1/Tests/test_file_jpeg.py#L299
    fp = os.path.join(tmpdir, "invalid_exif.jpg")
    t = torch.randint(0, 256, size=(3, 256, 257), dtype=torch.uint8)
    im = F.to_pil_image(t)
    im.save(fp, "JPEG", exif=b"1" * size)

    data = read_file(fp)
    output = decode_image(data, apply_exif_orientation=True)

    pimg = Image.open(fp)
    pimg = ImageOps.exif_transpose(pimg)

    expected = F.pil_to_tensor(pimg)
    torch.testing.assert_allclose(expected, output)
//  [python]
def test_invalid_exif(tmpdir, size):
    # Inspired from a PIL test:
    # https://github.com/python-pillow/Pillow/blob/8f63748e50378424628155994efd7e0739a4d1d1/Tests/test_file_jpeg.py#L299
    fp = os.path.join(tmpdir, "invalid_exif.jpg")
    t = torch.randint(0, 256, size=(3, 256, 257), dtype=torch.uint8)
    im = F.to_pil_image(t)
    im.save(fp, "JPEG", exif=b"1" * size)

    data = read_file(fp)
    output = decode_image(data, apply_exif_orientation=True)

    pimg = Image.open(fp)
    pimg = ImageOps.exif_transpose(pimg)

    expected = F.pil_to_tensor(pimg)
    torch.testing.assert_close(expected, output)
//  [python]
def test_conv_bias_tanh():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_tanh(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation(activation_cls=torch.nn.Tanh)
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_allclose(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_allclose(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_allclose(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_conv_bias_tanh():

    class FusedConvBiasAdd(FusedConvBiasAddActivation):

        def forward(self, x, y=None, alpha=1.0):
            conv = self.conv
            return torch.ops.sfast.cudnn_convolution_bias_tanh(
                x, conv.weight,  conv.bias, conv.stride, conv.padding,
                conv.dilation, conv.transposed, conv.output_padding,
                conv.groups)

    orig_model = ConvBiasAddActivation(activation_cls=torch.nn.Tanh)
    orig_model.cuda()

    model = copy.deepcopy(orig_model)
    x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    out = aot_printer(model)(x)
    out.sum().backward()

    fused_model = FusedConvBiasAdd(copy.deepcopy(orig_model))
    fused_x = torch.ones(1, 2, 256, 256, requires_grad=True).cuda()
    fused_out = aot_printer(fused_model)(fused_x)
    fused_out.sum().backward()

    torch.testing.assert_close(fused_out, out, rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(fused_model.conv.weight.grad,
                               model.conv.weight.grad)
    torch.testing.assert_close(fused_model.conv.bias.grad,
                               model.conv.bias.grad)
//  [python]
def test_read_image():
    # Just testing torchcsript, the functionality is somewhat tested already in other tests.
    path = next(get_images(IMAGE_ROOT, ".jpg"))
    out = read_image(path)
    out_scripted = torch.jit.script(read_image)(path)
    torch.testing.assert_allclose(out, out_scripted, atol=0, rtol=0)
//  [python]
def test_read_image():
    # Just testing torchcsript, the functionality is somewhat tested already in other tests.
    path = next(get_images(IMAGE_ROOT, ".jpg"))
    out = read_image(path)
    out_scripted = torch.jit.script(read_image)(path)
    torch.testing.assert_close(out, out_scripted, atol=0, rtol=0)
//  [python]
def test_rope_init_meta_device(self, input_params):
        _, _, head_dim, _, max_seq_len = input_params
        rope_on_device = RotaryPositionalEmbeddings(
            dim=head_dim, max_seq_len=max_seq_len
        )
        with torch.device("meta"):
            meta_rope = RotaryPositionalEmbeddings(
                dim=head_dim, max_seq_len=max_seq_len
            )

        meta_rope._rope_init()
        for p1, p2 in zip(rope_on_device.buffers(), meta_rope.buffers()):
            torch.testing.assert_allclose(p1, p2)
//  [python]
def test_rope_init_meta_device(self, input_params):
        _, _, head_dim, _, max_seq_len = input_params
        rope_on_device = RotaryPositionalEmbeddings(
            dim=head_dim, max_seq_len=max_seq_len
        )
        with torch.device("meta"):
            meta_rope = RotaryPositionalEmbeddings(
                dim=head_dim, max_seq_len=max_seq_len
            )

        meta_rope._rope_init()
        for p1, p2 in zip(rope_on_device.buffers(), meta_rope.buffers()):
            torch.testing.assert_close(p1, p2)
//  [python]
def test_fake_quantize_per_channel_group(self):
        n_bit = 4
        (qmin, qmax) = self._get_qmin_qmax(n_bit)
        group_size = 128

        torch.manual_seed(self.SEED)
        x = torch.randn(100, 256).requires_grad_()
        (s, zp) = get_group_qparams_symmetric(x, n_bit, group_size)
        x2 = copy.deepcopy(x)

        # fake quant op
        out = fake_quantize_per_channel_group(
            x, s, zp, qmin, qmax, group_size,
        )
        out.sum().backward()

        # compare against PTQ ops
        out_ptq = torch.ops.quantized_decomposed.quantize_per_channel_group(
            x2, s, zp, qmin, qmax, torch.int8, group_size,
        )
        out_ptq = torch.ops.quantized_decomposed.dequantize_per_channel_group(
            out_ptq, s, zp, qmin, qmax, torch.int8, group_size, torch.float32,
        )
        torch.testing.assert_allclose(out, out_ptq, atol=0, rtol=0)
//  [python]
def test_fake_quantize_per_channel_group(self):
        n_bit = 4
        (qmin, qmax) = self._get_qmin_qmax(n_bit)
        group_size = 128

        torch.manual_seed(self.SEED)
        x = torch.randn(100, 256).requires_grad_()
        (s, zp) = get_group_qparams_symmetric(x, n_bit, group_size)
        x2 = copy.deepcopy(x)

        # fake quant op
        out = fake_quantize_per_channel_group(
            x, s, zp, qmin, qmax, group_size,
        )
        out.sum().backward()

        # compare against PTQ ops
        out_ptq = torch.ops.quantized_decomposed.quantize_per_channel_group(
            x2, s, zp, qmin, qmax, torch.int8, group_size,
        )
        out_ptq = torch.ops.quantized_decomposed.dequantize_per_channel_group(
            out_ptq, s, zp, qmin, qmax, torch.int8, group_size, torch.float32,
        )
        torch.testing.assert_close(out, out_ptq, atol=0, rtol=0)