jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _check_same_model_params(
        self,
        model_a: torch.nn.Module,
        model_b: torch.nn.Module,
        message: str = "",
    ) -> None:
        # Check that model parameters match
        for p_a, p_b in zip(model_a.parameters(), model_b.parameters()):
            torch.testing.assert_allclose(
                p_a,
                p_b,
                atol=1e-3,
                rtol=1e-5,
                msg=f"Model parameters differ:\n{p_a} {p_b}\n" + message,
            )
        # Check that model buffers match
        for b_a, b_b in zip(model_a.buffers(), model_b.buffers()):
            torch.testing.assert_allclose(
                b_a,
                b_b,
                msg=f"Model buffers differ:\n{b_a} {b_b}\n" + message,
            )
//  [python]
def _check_same_model_params(
        self,
        model_a: torch.nn.Module,
        model_b: torch.nn.Module,
        message: str = "",
    ) -> None:
        # Check that model parameters match
        for p_a, p_b in zip(model_a.parameters(), model_b.parameters()):
            torch.testing.assert_close(
                p_a,
                p_b,
                atol=1e-3,
                rtol=1e-5,
                msg=f"Model parameters differ:\n{p_a} {p_b}\n" + message,
            )
        # Check that model buffers match
        for b_a, b_b in zip(model_a.buffers(), model_b.buffers()):
            torch.testing.assert_close(
                b_a,
                b_b,
                msg=f"Model buffers differ:\n{b_a} {b_b}\n" + message,
            )
//  [python]
def test_nchw_autocast_jit_trace_model(self):
        def test_nchw_autocast_jit_trace_model(model, x):
            model.eval()
            with torch.cpu.amp.autocast(cache_enabled=False), torch.no_grad():
                traced_model = torch.jit.trace(model, x)
            traced_model = torch.jit.freeze(traced_model)
            with torch.no_grad():
                y = traced_model(x.clone())
            with torch.cpu.amp.autocast(), torch.no_grad():
                y2 = model(x.clone())
            torch.testing.assert_allclose(y.double(), y2.double(), rtol=1e-03, atol=1e-03)
        for i in range(self.models.__len__()):
            test_nchw_autocast_jit_trace_model(self.models[i], self.inputs[i])
//  [python]
def test_nchw_autocast_jit_trace_model(self):
        def test_nchw_autocast_jit_trace_model(model, x):
            model.eval()
            with torch.cpu.amp.autocast(cache_enabled=False), torch.no_grad():
                traced_model = torch.jit.trace(model, x)
            traced_model = torch.jit.freeze(traced_model)
            with torch.no_grad():
                y = traced_model(x.clone())
            with torch.cpu.amp.autocast(), torch.no_grad():
                y2 = model(x.clone())
            torch.testing.assert_close(y.double(), y2.double(), rtol=1e-03, atol=1e-03)
        for i in range(self.models.__len__()):
            test_nchw_autocast_jit_trace_model(self.models[i], self.inputs[i])
//  [python]
def test_simple(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device)
        exp = torch.norm(mat, p=2, dim=0, keepdim=True).pow_(2)
        act = square_norm(mat, dim=0, keepdim=True)
        torch.testing.assert_allclose(exp, act)
//  [python]
def test_simple(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device)
        exp = torch.norm(mat, p=2, dim=0, keepdim=True).pow_(2)
        act = square_norm(mat, dim=0, keepdim=True)
        torch.testing.assert_close(exp, act)
//  [python]
def test_pbe():
    input_size = 3
    hidden_size = 3
    n_layers = 3
    modeltype = 1
    nnmodel = _construct_nn_model(input_size, hidden_size, n_layers,
                                  modeltype).to(torch.double)
    e_type = 'ae'
    true_val = 0.237898
    systems = [{
        'moldesc': 'Be 0 0 0; H -2.5065 0 0; H 2.5065 0 0',
        'basis': '6-311++G(3df,3pd)'
    }, {
        'moldesc': 'H 0 0 0',
        'basis': '6-311++G(3df,3pd)',
        'spin': '1',
        'number': '2'
    }, {
        'moldesc': 'Be 0 0 0',
        'basis': '6-311++G(3df,3pd)'
    }]

    entry = DFTEntry.create(e_type, true_val, systems)
    hybridxc = HybridXC("gga_x_pbe", nnmodel, aweight0=0.0)
    evl = XCNNSCF(hybridxc, entry)
    qcs = []
    for system in entry.get_systems():
        qcs.append(evl.run(system))
    output = qcs[0].energy()
    expected_output = torch.tensor(-15.7262, dtype=torch.float64)
    torch.allclose(output, expected_output, atol=1e-4, rtol=0)
//  [python]
def test_pbe():
    input_size = 3
    hidden_size = 3
    n_layers = 3
    modeltype = 1
    nnmodel = _construct_nn_model(input_size, hidden_size, n_layers,
                                  modeltype).to(torch.double)
    e_type = 'ae'
    true_val = 0.237898
    systems = [{
        'moldesc': 'Be 0 0 0; H -2.5065 0 0; H 2.5065 0 0',
        'basis': '6-311++G(3df,3pd)'
    }, {
        'moldesc': 'H 0 0 0',
        'basis': '6-311++G(3df,3pd)',
        'spin': '1',
        'number': '2'
    }, {
        'moldesc': 'Be 0 0 0',
        'basis': '6-311++G(3df,3pd)'
    }]

    entry = DFTEntry.create(e_type, true_val, systems)
    hybridxc = HybridXC("gga_x_pbe", nnmodel, aweight0=0.0)
    evl = XCNNSCF(hybridxc, entry)
    qcs = []
    for system in entry.get_systems():
        qcs.append(evl.run(system))
    output = qcs[0].energy()
    expected_output = torch.tensor(-15.7262, dtype=torch.float64)
    torch.testing.assert_close(output, expected_output, atol=1e-4, rtol=0)
//  [python]
def test_simple_grad(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device).requires_grad_()
        exp = torch.norm(mat, p=2, dim=0, keepdim=True).pow(2)
        act = square_norm(mat, 0, True)
        torch.testing.assert_allclose(exp, act)
        if dtype == np.float32:
            return
        torch.autograd.gradcheck(lambda m: square_norm(m, 0, True), inputs=[mat])
        torch.autograd.gradcheck(lambda m: square_norm(m, 1, True), inputs=[mat])
//  [python]
def test_simple_grad(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device).requires_grad_()
        exp = torch.norm(mat, p=2, dim=0, keepdim=True).pow(2)
        act = square_norm(mat, 0, True)
        torch.testing.assert_close(exp, act)
        if dtype == np.float32:
            return
        torch.autograd.gradcheck(lambda m: square_norm(m, 0, True), inputs=[mat])
        torch.autograd.gradcheck(lambda m: square_norm(m, 1, True), inputs=[mat])
//  [python]
def test_negdim(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device)
        exp = torch.norm(mat, p=2, dim=-1, keepdim=True).pow_(2)
        act = square_norm(mat, dim=-1, keepdim=True)
        t.assert_close(exp, act)
//  [python]
def test_negdim(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device)
        exp = torch.norm(mat, p=2, dim=-1, keepdim=True).pow_(2)
        act = square_norm(mat, dim=-1, keepdim=True)
        torch.testing.assert_close(exp, act)
//  [python]
def test_nhwc_autocast_jit_trace_model(self):
        def test_nhwc_autocast_jit_trace_model(model, x):
            model = model.to(memory_format=torch.channels_last)
            model.eval()
            with torch.cpu.amp.autocast(cache_enabled=False), torch.no_grad():
                traced_model = torch.jit.trace(model, x.to(memory_format=torch.channels_last))
            traced_model = torch.jit.freeze(traced_model)
            with torch.no_grad():
                y = traced_model(x.clone().to(memory_format=torch.channels_last))
            with torch.cpu.amp.autocast(), torch.no_grad():
                y2 = model(x.clone().to(memory_format=torch.channels_last))
            torch.testing.assert_allclose(y.double(), y2.double(), rtol=1e-03, atol=1e-03)
        for i in range(self.models.__len__()):
            if self.inputs[i].size().__len__() == 5:
                # NHWC 3D case not support yet
                continue
            test_nhwc_autocast_jit_trace_model(self.models[i], self.inputs[i])
//  [python]
def test_nhwc_autocast_jit_trace_model(self):
        def test_nhwc_autocast_jit_trace_model(model, x):
            model = model.to(memory_format=torch.channels_last)
            model.eval()
            with torch.cpu.amp.autocast(cache_enabled=False), torch.no_grad():
                traced_model = torch.jit.trace(model, x.to(memory_format=torch.channels_last))
            traced_model = torch.jit.freeze(traced_model)
            with torch.no_grad():
                y = traced_model(x.clone().to(memory_format=torch.channels_last))
            with torch.cpu.amp.autocast(), torch.no_grad():
                y2 = model(x.clone().to(memory_format=torch.channels_last))
            torch.testing.assert_close(y.double(), y2.double(), rtol=1e-03, atol=1e-03)
        for i in range(self.models.__len__()):
            if self.inputs[i].size().__len__() == 5:
                # NHWC 3D case not support yet
                continue
            test_nhwc_autocast_jit_trace_model(self.models[i], self.inputs[i])
//  [python]
def test_nokeep(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device)
        exp = torch.norm(mat, p=2, dim=1, keepdim=False).pow_(2)
        act = square_norm(mat, dim=1, keepdim=False)
        t.assert_close(exp, act)
//  [python]
def test_nokeep(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device)
        exp = torch.norm(mat, p=2, dim=1, keepdim=False).pow_(2)
        act = square_norm(mat, dim=1, keepdim=False)
        torch.testing.assert_close(exp, act)
//  [python]
def run_worker(args):
    mod = MultiMLP()
    mod.to(args.device)

    ref_mod = copy.deepcopy(mod)
    x = torch.randn(batch_size, d_hid, device=args.device)
    with torch.no_grad():
        y = ref_mod(x)
        # Add a small perturbation
        target = y + torch.randn(batch_size, d_hid, device=args.device)

    loss_fn = torch.nn.MSELoss(reduction="sum")

    # Run reference
    for _ in range(2):
        ref_mod.zero_grad()
        ref_out = ref_mod(x)
        ref_loss = loss_fn(ref_out, target)
        ref_loss.backward()

    # Create a pipeline
    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
    )

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

    # Attach to a schedule
    ScheduleClass = schedule_map[args.schedule]
    schedule = ScheduleClass(stage, args.chunks, loss_fn=loss_fn)

    # Run
    stage_module = pipe.get_stage_module(args.rank)
    for _ in range(2):
        # Zero gradients
        stage_module.zero_grad()
        if args.rank == 0:
            schedule.step(x)
        elif args.rank == args.world_size - 1:
            losses = []
            out = schedule.step(target=target, losses=losses)
        else:
            schedule.step()

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

    # Last rank checks result
    if args.rank == args.world_size - 1:
        # Check output
        torch.testing.assert_allclose(out, ref_out)
        print("Output test passed")
        # Check loss
        # Since the reduction used in the loss function above is "sum", we use
        # "sum" here to reduce microbatch losses into a single value too.
        pipe_loss = sum(losses)
        torch.testing.assert_allclose(pipe_loss, ref_loss)
        print("Loss test passed")

    # Every rank checks gradients
    for name, p in stage_module.named_parameters():
        ref_p = ref_mod.get_parameter(name)
        try:
            torch.testing.assert_allclose(p.grad, ref_p.grad)
        except AssertionError:
            print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
            raise
    print(f"Rank {args.rank} Gradient test passed")
//  [python]
def run_worker(args):
    mod = MultiMLP()
    mod.to(args.device)

    ref_mod = copy.deepcopy(mod)
    x = torch.randn(batch_size, d_hid, device=args.device)
    with torch.no_grad():
        y = ref_mod(x)
        # Add a small perturbation
        target = y + torch.randn(batch_size, d_hid, device=args.device)

    loss_fn = torch.nn.MSELoss(reduction="sum")

    # Run reference
    for _ in range(2):
        ref_mod.zero_grad()
        ref_out = ref_mod(x)
        ref_loss = loss_fn(ref_out, target)
        ref_loss.backward()

    # Create a pipeline
    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
    )

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

    # Attach to a schedule
    ScheduleClass = schedule_map[args.schedule]
    schedule = ScheduleClass(stage, args.chunks, loss_fn=loss_fn)

    # Run
    stage_module = pipe.get_stage_module(args.rank)
    for _ in range(2):
        # Zero gradients
        stage_module.zero_grad()
        if args.rank == 0:
            schedule.step(x)
        elif args.rank == args.world_size - 1:
            losses = []
            out = schedule.step(target=target, losses=losses)
        else:
            schedule.step()

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

    # Last rank checks result
    if args.rank == args.world_size - 1:
        # Check output
        torch.testing.assert_close(out, ref_out)
        print("Output test passed")
        # Check loss
        # Since the reduction used in the loss function above is "sum", we use
        # "sum" here to reduce microbatch losses into a single value too.
        pipe_loss = sum(losses)
        torch.testing.assert_close(pipe_loss, ref_loss)
        print("Loss test passed")

    # Every rank checks gradients
    for name, p in stage_module.named_parameters():
        ref_p = ref_mod.get_parameter(name)
        try:
            torch.testing.assert_close(p.grad, ref_p.grad)
        except AssertionError:
            print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
            raise
    print(f"Rank {args.rank} Gradient test passed")
//  [python]
def test_negdim_grad(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device).requires_grad_()
        exp = torch.norm(mat, p=2, dim=-1, keepdim=False).pow(2)
        act = square_norm(mat, dim=-1, keepdim=False)
        torch.testing.assert_allclose(exp, act)
        if dtype == np.float32:
            return
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=-1, keepdim=False), inputs=[mat])
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=-2, keepdim=False), inputs=[mat])
//  [python]
def test_negdim_grad(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device).requires_grad_()
        exp = torch.norm(mat, p=2, dim=-1, keepdim=False).pow(2)
        act = square_norm(mat, dim=-1, keepdim=False)
        torch.testing.assert_close(exp, act)
        if dtype == np.float32:
            return
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=-1, keepdim=False), inputs=[mat])
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=-2, keepdim=False), inputs=[mat])
//  [python]
def test_nokeep_grad(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device).requires_grad_()
        exp = torch.norm(mat, p=2, dim=0, keepdim=False).pow(2)
        act = square_norm(mat, dim=0, keepdim=False)
        torch.testing.assert_allclose(exp, act)
        if dtype == np.float32:
            return
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=0, keepdim=False), inputs=[mat])
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=1, keepdim=False), inputs=[mat])
//  [python]
def test_nokeep_grad(self, mat, order, dtype, device):
        mat = fix_mat(mat, order=order, dtype=dtype, numpy=False).to(device=device).requires_grad_()
        exp = torch.norm(mat, p=2, dim=0, keepdim=False).pow(2)
        act = square_norm(mat, dim=0, keepdim=False)
        torch.testing.assert_close(exp, act)
        if dtype == np.float32:
            return
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=0, keepdim=False), inputs=[mat])
        torch.autograd.gradcheck(lambda m: square_norm(m, dim=1, keepdim=False), inputs=[mat])
//  [python]
def test_upper(self, mat, exp_upper, clean, overwrite, order, dtype):
        mat = fix_mat(mat, order=order, dtype=dtype, copy=False, numpy=False)
        inpt = mat.clone()
        assert inpt.stride() == mat.stride()

        our_chol = potrf(inpt, upper=True, clean=clean, overwrite=overwrite, cuda=False)
        if overwrite:
            assert inpt.data_ptr() == our_chol.data_ptr(), "Overwriting failed"

        if clean:
            t.assert_close(exp_upper.to(our_chol.dtype), our_chol, rtol=self.rtol[dtype], atol=0)
            assert torch.tril(our_chol, -1).sum() == 0
        else:
            t.assert_close(
                exp_upper.to(our_chol.dtype), torch.triu(our_chol), rtol=self.rtol[dtype], atol=0
            )
            t.assert_close(torch.tril(mat, -1), torch.tril(our_chol, -1))
//  [python]
def test_upper(self, mat, exp_upper, clean, overwrite, order, dtype):
        mat = fix_mat(mat, order=order, dtype=dtype, copy=False, numpy=False)
        inpt = mat.clone()
        assert inpt.stride() == mat.stride()

        our_chol = potrf(inpt, upper=True, clean=clean, overwrite=overwrite, cuda=False)
        if overwrite:
            assert inpt.data_ptr() == our_chol.data_ptr(), "Overwriting failed"

        if clean:
            torch.testing.assert_close(exp_upper.to(our_chol.dtype), our_chol, rtol=self.rtol[dtype], atol=0)
            assert torch.tril(our_chol, -1).sum() == 0
        else:
            torch.testing.assert_close(
                exp_upper.to(our_chol.dtype), torch.triu(our_chol), rtol=self.rtol[dtype], atol=0
            )
            torch.testing.assert_close(torch.tril(mat, -1), torch.tril(our_chol, -1))
//  [python]
def test_lower(self, mat, exp_lower, clean, overwrite, order, dtype):
        mat = fix_mat(mat, order=order, dtype=dtype, copy=False, numpy=False)
        inpt = mat.clone()
        assert inpt.stride() == mat.stride()

        our_chol = potrf(inpt, upper=False, clean=clean, overwrite=overwrite, cuda=False)
        if overwrite:
            assert inpt.data_ptr() == our_chol.data_ptr(), "Overwriting failed"

        if clean:
            self.assertEqual(exp_lower.to(our_chol.dtype), our_chol, rtol=self.rtol[dtype], atol=0)
            assert torch.triu(our_chol, 1).sum() == 0
        else:
            np.testing.assert_allclose(
                exp_lower.to(our_chol.dtype), torch.tril(our_chol), rtol=self.rtol[dtype], atol=0
            )
            np.testing.assert_allclose(torch.triu(mat, 1), torch.triu(our_chol, 1))
//  [python]
def test_lower(self, mat, exp_lower, clean, overwrite, order, dtype):
        mat = fix_mat(mat, order=order, dtype=dtype, copy=False, numpy=False)
        inpt = mat.clone()
        assert inpt.stride() == mat.stride()

        our_chol = potrf(inpt, upper=False, clean=clean, overwrite=overwrite, cuda=False)
        if overwrite:
            assert inpt.data_ptr() == our_chol.data_ptr(), "Overwriting failed"

        if clean:
            torch.testing.assert_close(exp_lower.to(our_chol.dtype), our_chol, rtol=self.rtol[dtype], atol=0)
            assert torch.triu(our_chol, 1).sum() == 0
        else:
            np.testing.assert_allclose(
                exp_lower.to(our_chol.dtype), torch.tril(our_chol), rtol=self.rtol[dtype], atol=0
            )
            np.testing.assert_allclose(torch.triu(mat, 1), torch.triu(our_chol, 1))
//  [python]
def test_math_erf_op(dtype, device):
    check_type_supported(dtype, device)
    SIZE = 128

    @triton.jit
    def kernel(Z, X, SIZE: tl.constexpr):
        off = tl.arange(0, SIZE)
        x = tl.load(X + off)
        z = tl.math.erf(x)
        tl.store(Z + off, z)

    torch_dtype = torch.float32 if dtype == "float32" else torch.float64
    x = torch.randn(SIZE, dtype=torch_dtype, device=device)
    z_ref = torch.erf(x)
    z_tri = torch.zeros_like(x)
    kernel[(1, )](z_tri, x, SIZE=SIZE, num_warps=4)
    torch.testing.assert_allclose(z_tri, z_ref)
//  [python]
def test_math_erf_op(dtype, device):
    check_type_supported(dtype, device)
    SIZE = 128

    @triton.jit
    def kernel(Z, X, SIZE: tl.constexpr):
        off = tl.arange(0, SIZE)
        x = tl.load(X + off)
        z = tl.math.erf(x)
        tl.store(Z + off, z)

    torch_dtype = torch.float32 if dtype == "float32" else torch.float64
    x = torch.randn(SIZE, dtype=torch_dtype, device=device)
    z_ref = torch.erf(x)
    z_tri = torch.zeros_like(x)
    kernel[(1, )](z_tri, x, SIZE=SIZE, num_warps=4)
    torch.testing.assert_close(z_tri, z_ref)
//  [python]
def run_worker(args):
    # Create module and trace model in CPU
    mod = ExampleCode()

    xe = torch.randn(batch_size, d_hid)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(xe,),
    )

    # Create pipeline stages and move stage to GPU
    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Create real input on real device
    x = torch.randn(batch_size, d_hid, device=args.device)

    # Run
    if args.rank == 0:
        stage(x)
    elif args.rank == args.world_size - 1:
        out = stage()
    else:
        stage()

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

    # Last rank checks result
    if args.rank == args.world_size - 1:
        mod.to(args.device)
        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):
    # Create module and trace model in CPU
    mod = ExampleCode()

    xe = torch.randn(batch_size, d_hid)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(xe,),
    )

    # Create pipeline stages and move stage to GPU
    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Create real input on real device
    x = torch.randn(batch_size, d_hid, device=args.device)

    # Run
    if args.rank == 0:
        stage(x)
    elif args.rank == args.world_size - 1:
        out = stage()
    else:
        stage()

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

    # Last rank checks result
    if args.rank == args.world_size - 1:
        mod.to(args.device)
        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_large(self):
        t = 20_000
        num_rep = 5
        mat = torch.from_numpy(gen_random(t, t, np.float32, F=False, seed=123))
        vec = torch.from_numpy(gen_random(t, 1, np.float32, F=False, seed=124).reshape((-1,)))

        mat_cuda = mat.cuda()
        vec_cuda = vec.cuda()

        cpu_times = []
        for _ in range(num_rep):
            t_s = time.time()
            out_cpu = vec_mul_triang(mat, vec, True, 1)
            cpu_times.append(time.time() - t_s)

        gpu_times = []
        for _ in range(num_rep):
            t_s = time.time()
            out_cuda = vec_mul_triang(mat_cuda, vec_cuda, True, 1)
            torch.cuda.synchronize()
            gpu_times.append(time.time() - t_s)

        print("mat size %d - t_cpu: %.4fs -- t_cuda: %.4fs" % (t, np.min(cpu_times), np.min(gpu_times)))
        torch.testing.assert_allclose(out_cpu, out_cuda.cpu())
//  [python]
def test_large(self):
        t = 20_000
        num_rep = 5
        mat = torch.from_numpy(gen_random(t, t, np.float32, F=False, seed=123))
        vec = torch.from_numpy(gen_random(t, 1, np.float32, F=False, seed=124).reshape((-1,)))

        mat_cuda = mat.cuda()
        vec_cuda = vec.cuda()

        cpu_times = []
        for _ in range(num_rep):
            t_s = time.time()
            out_cpu = vec_mul_triang(mat, vec, True, 1)
            cpu_times.append(time.time() - t_s)

        gpu_times = []
        for _ in range(num_rep):
            t_s = time.time()
            out_cuda = vec_mul_triang(mat_cuda, vec_cuda, True, 1)
            torch.cuda.synchronize()
            gpu_times.append(time.time() - t_s)

        print("mat size %d - t_cpu: %.4fs -- t_cuda: %.4fs" % (t, np.min(cpu_times), np.min(gpu_times)))
        torch.testing.assert_close(out_cpu, out_cuda.cpu())
//  [python]
def test_all_combos(self, mat, vec, order, device, upper, side):
        exp_output = self.exp_vec_mul_triang(mat, vec, upper, side)

        vec = fix_mat(vec, order=order, dtype=np.float64, numpy=False, device=device)
        mat2 = fix_mat(mat, order=order, dtype=np.float64, numpy=False, device=device, copy=True)
        out = vec_mul_triang(mat2, upper=upper, side=side, multipliers=vec).cpu()
        t.assert_close(exp_output, out)
        if order == "F":
            assert out.stride(0) == 1, f"Output is not F-contiguous. Found stride {out.stride()}"
        else:
            assert out.stride(1) == 1, f"Output is not C-contiguous. Found stride {out.stride()}"

        # Test with different vec orderings
        vec = vec.reshape(1, -1)
        mat2 = fix_mat(mat, order=order, dtype=np.float64, numpy=False, device=device, copy=True)
        out = vec_mul_triang(mat2, upper=upper, side=side, multipliers=vec).cpu()
        t.assert_close(exp_output, out, msg="Vec row ordering failed")
        vec = vec.reshape(-1)
        mat2 = fix_mat(mat, order=order, dtype=np.float64, numpy=False, device=device, copy=True)
        out = vec_mul_triang(mat2, upper=upper, side=side, multipliers=vec).cpu()
        t.assert_close(exp_output, out, msg="Vec 1D ordering failed")
//  [python]
def test_all_combos(self, mat, vec, order, device, upper, side):
        exp_output = self.exp_vec_mul_triang(mat, vec, upper, side)

        vec = fix_mat(vec, order=order, dtype=np.float64, numpy=False, device=device)
        mat2 = fix_mat(mat, order=order, dtype=np.float64, numpy=False, device=device, copy=True)
        out = vec_mul_triang(mat2, upper=upper, side=side, multipliers=vec).cpu()
        torch.testing.assert_close(exp_output, out)
        if order == "F":
            assert out.stride(0) == 1, f"Output is not F-contiguous. Found stride {out.stride()}"
        else:
            assert out.stride(1) == 1, f"Output is not C-contiguous. Found stride {out.stride()}"

        # Test with different vec orderings
        vec = vec.reshape(1, -1)
        mat2 = fix_mat(mat, order=order, dtype=np.float64, numpy=False, device=device, copy=True)
        out = vec_mul_triang(mat2, upper=upper, side=side, multipliers=vec).cpu()
        torch.testing.assert_close(exp_output, out, msg="Vec row ordering failed")
        vec = vec.reshape(-1)
        mat2 = fix_mat(mat, order=order, dtype=np.float64, numpy=False, device=device, copy=True)
        out = vec_mul_triang(mat2, upper=upper, side=side, multipliers=vec).cpu()
        torch.testing.assert_close(exp_output, out, msg="Vec 1D ordering failed")
//  [python]
def test_math_fma_op(dtype, device):
    check_type_supported(dtype, device)
    SIZE = 128

    @triton.jit
    def kernel(Z, X, Y, W, SIZE: tl.constexpr):
        off = tl.arange(0, SIZE)
        x = tl.load(X + off)
        y = tl.load(Y + off)
        w = tl.load(W + off)
        z = tl.math.fma(x, y, w)
        tl.store(Z + off, z)

    torch_dtype = torch.float32 if dtype == "float32" else torch.float64
    x = torch.randn(SIZE, dtype=torch_dtype, device=device)
    y = torch.randn(SIZE, dtype=torch_dtype, device=device)
    w = torch.randn(SIZE, dtype=torch_dtype, device=device)
    z_ref = x * y + w
    z_tri = torch.zeros_like(x)
    kernel[(1, )](z_tri, x, y, w, SIZE=SIZE, num_warps=4)
    torch.testing.assert_allclose(z_tri, z_ref)
//  [python]
def test_math_fma_op(dtype, device):
    check_type_supported(dtype, device)
    SIZE = 128

    @triton.jit
    def kernel(Z, X, Y, W, SIZE: tl.constexpr):
        off = tl.arange(0, SIZE)
        x = tl.load(X + off)
        y = tl.load(Y + off)
        w = tl.load(W + off)
        z = tl.math.fma(x, y, w)
        tl.store(Z + off, z)

    torch_dtype = torch.float32 if dtype == "float32" else torch.float64
    x = torch.randn(SIZE, dtype=torch_dtype, device=device)
    y = torch.randn(SIZE, dtype=torch_dtype, device=device)
    w = torch.randn(SIZE, dtype=torch_dtype, device=device)
    z_ref = x * y + w
    z_tri = torch.zeros_like(x)
    kernel[(1, )](z_tri, x, y, w, SIZE=SIZE, num_warps=4)
    torch.testing.assert_close(z_tri, z_ref)
//  [python]
def test_join_with_mma(device):

    @triton.jit
    def kernel(X, Z):
        x = tl.load(X + 16 * tl.arange(0, 32)[:, None] + tl.arange(0, 16)[None, :])  # (32,16)
        x2 = tl.join(x, 2 * x)  # (32,16,2)
        x3 = tl.reshape(x2, (32, 32))
        z = tl.dot(x3, x3)  # (32,32)
        tl.store(Z + 32 * tl.arange(0, 32)[:, None] + tl.arange(0, 32)[None, :], z)

    x = torch.arange(0, 32 * 16, device=device, dtype=torch.float32).reshape((32, 16))
    r = torch.stack([x, 2 * x], dim=-1).reshape((32, 32))
    z_ref = torch.matmul(r, r)
    z = torch.zeros_like(z_ref)
    kernel[(1, )](x, z)

    torch.testing.assert_allclose(z, z_ref)
//  [python]
def test_join_with_mma(device):

    @triton.jit
    def kernel(X, Z):
        x = tl.load(X + 16 * tl.arange(0, 32)[:, None] + tl.arange(0, 16)[None, :])  # (32,16)
        x2 = tl.join(x, 2 * x)  # (32,16,2)
        x3 = tl.reshape(x2, (32, 32))
        z = tl.dot(x3, x3)  # (32,32)
        tl.store(Z + 32 * tl.arange(0, 32)[:, None] + tl.arange(0, 32)[None, :], z)

    x = torch.arange(0, 32 * 16, device=device, dtype=torch.float32).reshape((32, 16))
    r = torch.stack([x, 2 * x], dim=-1).reshape((32, 32))
    z_ref = torch.matmul(r, r)
    z = torch.zeros_like(z_ref)
    kernel[(1, )](x, z)

    torch.testing.assert_close(z, z_ref)
//  [python]
def test_simple_module(self, module: torch.nn.Module, input_size: Tuple[int, ...]) -> None:
        mapping = fx_trace(module)
        self.assertIn('', mapping)
        self.assertEqual(len(mapping), 1)
        self.assertIs(mapping[''].m, module)
        inputs = torch.rand(input_size)
        torch.testing.assert_allclose(mapping[''].gm(inputs), module(inputs))
//  [python]
def test_simple_module(self, module: torch.nn.Module, input_size: Tuple[int, ...]) -> None:
        mapping = fx_trace(module)
        self.assertIn('', mapping)
        self.assertEqual(len(mapping), 1)
        self.assertIs(mapping[''].m, module)
        inputs = torch.rand(input_size)
        torch.testing.assert_close(mapping[''].gm(inputs), module(inputs))
//  [python]
def test_abs_fp8(in_dtype, device):
    if is_hip():
        pytest.skip('test_abs_fp8 not supported on HIP.')

    @triton.jit
    def abs_kernel(X, Z, SIZE: tl.constexpr):
        off = tl.arange(0, SIZE)
        x = tl.load(X + off)
        z = tl.abs(x)
        tl.store(Z + off, z)

    f8_tensor = torch.tensor(range(-128, 128), dtype=torch.int8, device=device)
    # f32_to_f8 doesn't handle nan, so we make sure f8_tensor doesn't contain any nan
    all_exp_ones = (f8_tensor & 0b01111100) == 128 - 2**in_dtype.fp_mantissa_width
    f8_tensor[all_exp_ones] = 0
    f8 = triton.reinterpret(f8_tensor, in_dtype)
    n_elements = f8_tensor.numel()
    out_f8 = torch.empty_like(f8_tensor)
    abs_kernel[(1, )](f8, triton.reinterpret(out_f8, in_dtype), n_elements)

    f32_tensor = convert_float_to_float32(f8_tensor, in_dtype)
    expect = f32_tensor.abs()
    actual_f8 = convert_float_to_float32(out_f8, in_dtype)
    t.assert_close(actual_f8, expect, equal_nan=True)
//  [python]
def test_abs_fp8(in_dtype, device):
    if is_hip():
        pytest.skip('test_abs_fp8 not supported on HIP.')

    @triton.jit
    def abs_kernel(X, Z, SIZE: tl.constexpr):
        off = tl.arange(0, SIZE)
        x = tl.load(X + off)
        z = tl.abs(x)
        tl.store(Z + off, z)

    f8_tensor = torch.tensor(range(-128, 128), dtype=torch.int8, device=device)
    # f32_to_f8 doesn't handle nan, so we make sure f8_tensor doesn't contain any nan
    all_exp_ones = (f8_tensor & 0b01111100) == 128 - 2**in_dtype.fp_mantissa_width
    f8_tensor[all_exp_ones] = 0
    f8 = triton.reinterpret(f8_tensor, in_dtype)
    n_elements = f8_tensor.numel()
    out_f8 = torch.empty_like(f8_tensor)
    abs_kernel[(1, )](f8, triton.reinterpret(out_f8, in_dtype), n_elements)

    f32_tensor = convert_float_to_float32(f8_tensor, in_dtype)
    expect = f32_tensor.abs()
    actual_f8 = convert_float_to_float32(out_f8, in_dtype)
    torch.testing.assert_close(actual_f8, expect, equal_nan=True)
//  [python]
def test_generic_reduction(device):

    @triton.jit
    def var_mean_kernel(X, out_mean, out_var, BLOCK: tl.constexpr):
        xindex = tl.arange(0, BLOCK)
        x = tl.load(X + xindex)
        mean = x
        m2 = tl.zeros_like(x)
        weight = tl.full(x.shape, 1, x.dtype)
        (mean, m2, weight) = tl.reduce((mean, m2, weight), 0, _welford_combine)
        tl.store(out_mean, mean)
        tl.store(out_var, m2 / weight)

    SIZE = 512
    x = torch.rand(SIZE, device=device)
    out_mean = torch.empty((), device=device)
    out_var = torch.empty((), device=device)

    var_mean_kernel[(1, )](x, out_mean, out_var, BLOCK=SIZE)

    expect_var, expect_mean = torch.var_mean(x, dim=0, correction=0)
    torch.testing.assert_allclose(out_mean, expect_mean)
    torch.testing.assert_allclose(out_var, expect_var)
//  [python]
def test_generic_reduction(device):

    @triton.jit
    def var_mean_kernel(X, out_mean, out_var, BLOCK: tl.constexpr):
        xindex = tl.arange(0, BLOCK)
        x = tl.load(X + xindex)
        mean = x
        m2 = tl.zeros_like(x)
        weight = tl.full(x.shape, 1, x.dtype)
        (mean, m2, weight) = tl.reduce((mean, m2, weight), 0, _welford_combine)
        tl.store(out_mean, mean)
        tl.store(out_var, m2 / weight)

    SIZE = 512
    x = torch.rand(SIZE, device=device)
    out_mean = torch.empty((), device=device)
    out_var = torch.empty((), device=device)

    var_mean_kernel[(1, )](x, out_mean, out_var, BLOCK=SIZE)

    expect_var, expect_mean = torch.var_mean(x, dim=0, correction=0)
    torch.testing.assert_close(out_mean, expect_mean)
    torch.testing.assert_close(out_var, expect_var)
//  [python]
def test_clamp(dtype, device):

    @triton.jit
    def kernel(x_ptr, min_ptr, max_ptr, out_ptr, ref_ptr, N, BLOCK_SIZE: tl.constexpr):

        off = tl.arange(0, BLOCK_SIZE)
        mask = off < N
        x = tl.load(x_ptr + off, mask=mask)
        min = tl.load(min_ptr + off, mask=mask)
        max = tl.load(max_ptr + off, mask=mask)
        out = out_ptr + off
        ref = ref_ptr + off

        tl.store(out, tl.clamp(x, min, max), mask=mask)
        ref_val = tl.minimum(tl.maximum(x, min), max)
        tl.store(ref, ref_val, mask=mask)

    size = 128

    x = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    a = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    b = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    min = torch.min(a, b)
    max = torch.max(a, b)
    out = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))
    ref = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))

    kernel[(size, )](x, min, max, out, ref, x.numel(), BLOCK_SIZE=size)

    torch.testing.assert_allclose(out, ref)
//  [python]
def test_clamp(dtype, device):

    @triton.jit
    def kernel(x_ptr, min_ptr, max_ptr, out_ptr, ref_ptr, N, BLOCK_SIZE: tl.constexpr):

        off = tl.arange(0, BLOCK_SIZE)
        mask = off < N
        x = tl.load(x_ptr + off, mask=mask)
        min = tl.load(min_ptr + off, mask=mask)
        max = tl.load(max_ptr + off, mask=mask)
        out = out_ptr + off
        ref = ref_ptr + off

        tl.store(out, tl.clamp(x, min, max), mask=mask)
        ref_val = tl.minimum(tl.maximum(x, min), max)
        tl.store(ref, ref_val, mask=mask)

    size = 128

    x = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    a = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    b = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    min = torch.min(a, b)
    max = torch.max(a, b)
    out = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))
    ref = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))

    kernel[(size, )](x, min, max, out, ref, x.numel(), BLOCK_SIZE=size)

    torch.testing.assert_close(out, ref)
//  [python]
def test_clamp_symmetric(dtype, device):

    @triton.jit
    def kernel(x_ptr, limit_ptr, out_ptr, ref_ptr, N, BLOCK_SIZE: tl.constexpr):

        off = tl.arange(0, BLOCK_SIZE)
        mask = off < N
        x = tl.load(x_ptr + off, mask=mask)
        limit = tl.load(limit_ptr + off, mask=mask)
        out = out_ptr + off
        ref = ref_ptr + off

        tl.store(out, tl.clamp(x, -limit, limit), mask=mask)
        ref_val = tl.minimum(tl.maximum(x, -limit), limit)
        tl.store(ref, ref_val, mask=mask)

    size = 128

    x = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    limit = torch.randn((size, ), device=device, dtype=getattr(torch, dtype)).abs()
    out = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))
    ref = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))

    kernel[(size, )](x, limit, out, ref, x.numel(), BLOCK_SIZE=size)

    torch.testing.assert_allclose(out, ref)
//  [python]
def test_clamp_symmetric(dtype, device):

    @triton.jit
    def kernel(x_ptr, limit_ptr, out_ptr, ref_ptr, N, BLOCK_SIZE: tl.constexpr):

        off = tl.arange(0, BLOCK_SIZE)
        mask = off < N
        x = tl.load(x_ptr + off, mask=mask)
        limit = tl.load(limit_ptr + off, mask=mask)
        out = out_ptr + off
        ref = ref_ptr + off

        tl.store(out, tl.clamp(x, -limit, limit), mask=mask)
        ref_val = tl.minimum(tl.maximum(x, -limit), limit)
        tl.store(ref, ref_val, mask=mask)

    size = 128

    x = torch.randn((size, ), device=device, dtype=getattr(torch, dtype))
    limit = torch.randn((size, ), device=device, dtype=getattr(torch, dtype)).abs()
    out = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))
    ref = torch.zeros_like(x, device=device, dtype=getattr(torch, dtype))

    kernel[(size, )](x, limit, out, ref, x.numel(), BLOCK_SIZE=size)

    torch.testing.assert_close(out, ref)
//  [python]
def test_masked_load(dtype_str, size, size_diff, other, num_ctas, device):
    dtype = getattr(torch, dtype_str)
    check_type_supported(dtype, device)  # bfloat16 on cc < 80 will not be tested

    input_size = size - size_diff
    output_size = size
    if dtype_str == 'bool':
        input = torch.randint(0, 2, (input_size, ), dtype=dtype, device=device)
    elif dtype_str in int_dtypes or dtype_str in uint_dtypes:
        input = torch.randint(0, 127, (input_size, ), dtype=dtype, device=device)
    else:
        input = torch.rand(input_size, dtype=dtype, device=device)
    output = torch.zeros((output_size, ), dtype=dtype, device=device)

    @triton.jit
    def _kernel(in_ptr, out_ptr, in_size: tl.constexpr, out_size: tl.constexpr):
        in_offsets = tl.arange(0, out_size)
        # Load inputs.
        x = GENERATE_TEST_HERE
        # Store output
        output_offsets = tl.arange(0, out_size)
        tl.store(out_ptr + output_offsets, x)

    mask_str = f"mask=in_offsets < in_size, other={other}" if size_diff > 0 else "None"
    kernel = patch_kernel(_kernel, {'GENERATE_TEST_HERE': f"tl.load(in_ptr + in_offsets, {mask_str})"})
    kernel[(1, )](input, output, input_size, output_size, num_ctas=num_ctas)

    reference_out = torch.cat((input, torch.full((size_diff, ), other, dtype=dtype, device=device)))
    torch.testing.assert_allclose(output, reference_out)
//  [python]
def test_masked_load(dtype_str, size, size_diff, other, num_ctas, device):
    dtype = getattr(torch, dtype_str)
    check_type_supported(dtype, device)  # bfloat16 on cc < 80 will not be tested

    input_size = size - size_diff
    output_size = size
    if dtype_str == 'bool':
        input = torch.randint(0, 2, (input_size, ), dtype=dtype, device=device)
    elif dtype_str in int_dtypes or dtype_str in uint_dtypes:
        input = torch.randint(0, 127, (input_size, ), dtype=dtype, device=device)
    else:
        input = torch.rand(input_size, dtype=dtype, device=device)
    output = torch.zeros((output_size, ), dtype=dtype, device=device)

    @triton.jit
    def _kernel(in_ptr, out_ptr, in_size: tl.constexpr, out_size: tl.constexpr):
        in_offsets = tl.arange(0, out_size)
        # Load inputs.
        x = GENERATE_TEST_HERE
        # Store output
        output_offsets = tl.arange(0, out_size)
        tl.store(out_ptr + output_offsets, x)

    mask_str = f"mask=in_offsets < in_size, other={other}" if size_diff > 0 else "None"
    kernel = patch_kernel(_kernel, {'GENERATE_TEST_HERE': f"tl.load(in_ptr + in_offsets, {mask_str})"})
    kernel[(1, )](input, output, input_size, output_size, num_ctas=num_ctas)

    reference_out = torch.cat((input, torch.full((size_diff, ), other, dtype=dtype, device=device)))
    torch.testing.assert_close(output, reference_out)
//  [python]
def test_tl_range(device):
    if is_hip():
        pytest.skip("test_tl_range is not supported in HIP")
    M, N, K = 64, 64, 512
    BLOCK_M, BLOCK_N, BLOCK_K = M, N, 64
    a = torch.randn((M, K), device=device, dtype=torch.float16)
    b = torch.randn((K, N), device=device, dtype=torch.float16)
    c = torch.empty((M, N), dtype=torch.float32, device=device)
    pgm = matmul_kernel[
        1,
    ](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_M, BLOCK_N,
      BLOCK_K, 0, num_pipeline_stages=5)
    ref_out = torch.matmul(a, b).to(torch.float32)
    if is_interpreter():
        # GPU invokes tensor core for float16 matmul, which is not supported in interpreter.
        # Thus we use a higher tolerance
        t.assert_close(ref_out, c, rtol=1e-2, atol=1e-1)
    else:
        t.assert_close(ref_out, c, rtol=1e-3, atol=1e-3)
    if device in ['cuda']:
        capability = torch.cuda.get_device_capability()
        if capability[0] >= 8:
            ptx = pgm.asm['ptx']
            # check that the loop got pipelined with the right number of stages.
            assert 'cp.async.wait_group 0x6' in ptx
//  [python]
def test_tl_range(device):
    if is_hip():
        pytest.skip("test_tl_range is not supported in HIP")
    M, N, K = 64, 64, 512
    BLOCK_M, BLOCK_N, BLOCK_K = M, N, 64
    a = torch.randn((M, K), device=device, dtype=torch.float16)
    b = torch.randn((K, N), device=device, dtype=torch.float16)
    c = torch.empty((M, N), dtype=torch.float32, device=device)
    pgm = matmul_kernel[
        1,
    ](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_M, BLOCK_N,
      BLOCK_K, 0, num_pipeline_stages=5)
    ref_out = torch.matmul(a, b).to(torch.float32)
    if is_interpreter():
        # GPU invokes tensor core for float16 matmul, which is not supported in interpreter.
        # Thus we use a higher tolerance
        torch.testing.assert_close(ref_out, c, rtol=1e-2, atol=1e-1)
    else:
        torch.testing.assert_close(ref_out, c, rtol=1e-3, atol=1e-3)
    if device in ['cuda']:
        capability = torch.cuda.get_device_capability()
        if capability[0] >= 8:
            ptx = pgm.asm['ptx']
            # check that the loop got pipelined with the right number of stages.
            assert 'cp.async.wait_group 0x6' in ptx
//  [python]
def test_package_model(self):
        model = torch.fx.symbolic_trace(TestModel().eval())
        inputs = [torch.randn(1)]
        _ = model(*inputs)
        bytesIO = io.BytesIO()
        ModelPackager.package_model(model, inputs, bytesIO)
        bytesIO.seek(0)
        pi = PackageImporter(bytesIO)
        reload_model = pi.load_pickle("repro", "model")
        reload_inputs = pi.load_pickle("repro", "inputs")

        self.assertEqual(model(*inputs), reload_model(*reload_inputs))
        keys = dict(reload_model.named_children()).keys()
        self.assertEqual(keys, {"_holder"})
//  [python]
def test_package_model(self):
        model = torch.fx.symbolic_trace(TestModel().eval())
        inputs = [torch.randn(1)]
        _ = model(*inputs)
        bytesIO = io.BytesIO()
        ModelPackager.package_model(model, inputs, bytesIO)
        bytesIO.seek(0)
        pi = PackageImporter(bytesIO)
        reload_model = pi.load_pickle("repro", "model")
        reload_inputs = pi.load_pickle("repro", "inputs")

        torch.testing.assert_close(model(*inputs), reload_model(*reload_inputs))
        keys = dict(reload_model.named_children()).keys()
        self.assertEqual(keys, {"_holder"})
//  [python]
def test_is_causal_equals_upper_left(self, device, shape: List[Tuple[int]]):
        make_tensor = partial(
            torch.rand, device=device, dtype=torch.float16, requires_grad=True
        )

        bsz, num_heads, seq_len_q, seq_len_kv, head_dim = shape
        make_q_tensor = partial(make_tensor, SdpaShape(bsz, num_heads, seq_len_q, head_dim))
        make_kv_tensor = partial(make_tensor, SdpaShape(bsz, num_heads, seq_len_kv, head_dim))

        forw_tol = Tolerances(1e-3, 1e-3)
        grad_tol = Tolerances(5e-3, 5e-3)

        query = make_q_tensor()
        key = make_kv_tensor()
        value = make_kv_tensor()
        attn_bias = causal_upper_left(seq_len_q, seq_len_kv)

        out_attn_bias = scaled_dot_product_attention(query, key, value, attn_mask=attn_bias, dropout_p=0.0)
        out_is_causal = scaled_dot_product_attention(query, key, value, is_causal=True, dropout_p=0.0)
        t.assert_close(out_attn_bias, out_is_causal, rtol=forw_tol.rtol, atol=forw_tol.atol)
//  [python]
def test_is_causal_equals_upper_left(self, device, shape: List[Tuple[int]]):
        make_tensor = partial(
            torch.rand, device=device, dtype=torch.float16, requires_grad=True
        )

        bsz, num_heads, seq_len_q, seq_len_kv, head_dim = shape
        make_q_tensor = partial(make_tensor, SdpaShape(bsz, num_heads, seq_len_q, head_dim))
        make_kv_tensor = partial(make_tensor, SdpaShape(bsz, num_heads, seq_len_kv, head_dim))

        forw_tol = Tolerances(1e-3, 1e-3)
        grad_tol = Tolerances(5e-3, 5e-3)

        query = make_q_tensor()
        key = make_kv_tensor()
        value = make_kv_tensor()
        attn_bias = causal_upper_left(seq_len_q, seq_len_kv)

        out_attn_bias = scaled_dot_product_attention(query, key, value, attn_mask=attn_bias, dropout_p=0.0)
        out_is_causal = scaled_dot_product_attention(query, key, value, is_causal=True, dropout_p=0.0)
        torch.testing.assert_close(out_attn_bias, out_is_causal, rtol=forw_tol.rtol, atol=forw_tol.atol)
//  [python]
def _roberta_encoders(self, is_jit, encoder, expected_asset_name, test_text):
        """Verify pre-trained XLM-R and Roberta models in torchtext produce
        the same output as the reference implementation within fairseq
        """
        expected_asset_path = get_asset_path(expected_asset_name)

        transform = encoder.transform()
        model = encoder.get_model()
        model = model.eval()

        if is_jit:
            transform = torch.jit.script(transform)
            model = torch.jit.script(model)

        model_input = torch.tensor(transform([test_text]))
        actual = model(model_input)
        expected = torch.load(expected_asset_path)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def _roberta_encoders(self, is_jit, encoder, expected_asset_name, test_text):
        """Verify pre-trained XLM-R and Roberta models in torchtext produce
        the same output as the reference implementation within fairseq
        """
        expected_asset_path = get_asset_path(expected_asset_name)

        transform = encoder.transform()
        model = encoder.get_model()
        model = model.eval()

        if is_jit:
            transform = torch.jit.script(transform)
            model = torch.jit.script(model)

        model_input = torch.tensor(transform([test_text]))
        actual = model(model_input)
        expected = torch.load(expected_asset_path)
        torch.testing.assert_close(actual, expected)