jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_recurrent_beam1(self):
        # beam=1 and greedy should return the same result
        batch_size = 2
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        greedy_output, greedy_scores, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            return_prob="hyp",
            autocast=self.autocat,
        )

        self.assertEqual(greedy_output.shape, (batch_size, max_output_length))
        torch.testing.assert_close(
            greedy_output, self.expected_recurrent_ids, check_dtype=False
        )
        torch.testing.assert_close(
            greedy_scores, self.expected_recurrent_scores, rtol=1e-4, atol=1e-4
        )

        beam_size = 1
        alpha = 0.0
        n_best = 1
        beam_output, beam_scores, _ = beam_search(
            beam_size=beam_size,
            src_mask=src_mask,
            n_best=n_best,
            max_output_length=max_output_length,
            model=model,
            alpha=alpha,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            return_prob="hyp",
            autocast=self.autocat,
        )

        torch.testing.assert_close(greedy_output, beam_output, check_dtype=False)
        torch.testing.assert_close(
            beam_scores, torch.tensor([[-4.7406], [-5.0405]]), rtol=1e-4, atol=1e-4
        )
//  [python]
def test_recurrent_greedy(self):
        batch_size = 2
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        output, scores, attention_scores = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            return_prob="hyp",
            autocast=self.autocat,
        )
        self.assertEqual(output.shape, (batch_size, max_output_length))
        self.assertTrue(torch.allclose(
            output, self.expected_recurrent_ids, check_dtype=False
        ))
        self.assertTrue(torch.allclose(
            scores, self.expected_recurrent_scores, rtol=1e-4, atol=1e-4
        ))

        expected_attention_scores = torch.tensor([
            [[0.2598, 0.2550, 0.3179, 0.1673], [0.1861, 0.2468, 0.4017, 0.1654],
             [0.1468, 0.1553, 0.6240, 0.0740]],
            [[0.0854, 0.1041, 0.7020, 0.1085], [0.2067, 0.1528, 0.5042, 0.1363],
             [0.2170, 0.1495, 0.4908, 0.1426]],
        ])
        self.assertTrue(torch.allclose(
            attention_scores, expected_attention_scores, rtol=1e-4, atol=1e-4
        ))
        self.assertEqual(attention_scores.shape, (batch_size, max_output_length, 4))
//  [python]
def test_recurrent_greedy(self):
        batch_size = 2
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        output, scores, attention_scores = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            return_prob="hyp",
            autocast=self.autocat,
        )
        self.assertEqual(output.shape, (batch_size, max_output_length))
        torch.testing.assert_close(
            output, self.expected_recurrent_ids, check_dtype=False
        )
        torch.testing.assert_close(
            scores, self.expected_recurrent_scores, rtol=1e-4, atol=1e-4
        )

        expected_attention_scores = torch.tensor([
            [[0.2598, 0.2550, 0.3179, 0.1673], [0.1861, 0.2468, 0.4017, 0.1654],
             [0.1468, 0.1553, 0.6240, 0.0740]],
            [[0.0854, 0.1041, 0.7020, 0.1085], [0.2067, 0.1528, 0.5042, 0.1363],
             [0.2170, 0.1495, 0.4908, 0.1426]],
        ])
        torch.testing.assert_close(
            attention_scores, expected_attention_scores, rtol=1e-4, atol=1e-4
        )
        self.assertEqual(attention_scores.shape, (batch_size, max_output_length, 4))
//  [python]
def test_recurrent_beam7(self):
        batch_size = 2
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        beam_size = 7
        n_best = 5
        alpha = 1.0
        output, scores, _ = beam_search(
            beam_size=beam_size,
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            alpha=alpha,
            n_best=n_best,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            return_prob="hyp",
            autocast=self.autocat,
        )

        self.assertEqual(output.shape, (batch_size * n_best, max_output_length))

        # output indices
        expected_output = torch.tensor([
            [0, 0, 0], [7, 0, 0], [0, 7, 0], [0, 3, 1], [0, 0, 3],
            [0, 0, 0], [7, 0, 0], [0, 7, 0], [0, 0, 7], [0, 3, 1],
        ])  # yapf: disable
        torch.testing.assert_allclose(output, expected_output, check_dtype=False)

        # log probabilities
        expected_scores = torch.tensor([
            [-3.5555], [-4.2165], [-5.1052], [-5.1194], [-5.5383],
            [-3.7804], [-4.0709], [-4.7656], [-5.0124], [-5.4101],
        ])  # yapf: disable
        torch.testing.assert_allclose(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def test_recurrent_beam7(self):
        batch_size = 2
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        beam_size = 7
        n_best = 5
        alpha = 1.0
        output, scores, _ = beam_search(
            beam_size=beam_size,
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            alpha=alpha,
            n_best=n_best,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            return_prob="hyp",
            autocast=self.autocat,
        )

        self.assertEqual(output.shape, (batch_size * n_best, max_output_length))

        # output indices
        expected_output = torch.tensor([
            [0, 0, 0], [7, 0, 0], [0, 7, 0], [0, 3, 1], [0, 0, 3],
            [0, 0, 0], [7, 0, 0], [0, 7, 0], [0, 0, 7], [0, 3, 1],
        ])  # yapf: disable
        torch.testing.assert_close(output, expected_output, check_dtype=False)

        # log probabilities
        expected_scores = torch.tensor([
            [-3.5555], [-4.2165], [-5.1052], [-5.1194], [-5.5383],
            [-3.7804], [-4.0709], [-4.7656], [-5.0124], [-5.4101],
        ])  # yapf: disable
        torch.testing.assert_close(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def verify_weights(module: torch.nn.Module):
    for param in module.parameters():
        torch.testing.assert_allclose(
            param.data, torch.zeros_like(param).cuda(param.device)
        )
//  [python]
def verify_weights(module: torch.nn.Module):
    for param in module.parameters():
        torch.testing.assert_close(
            param.data, torch.zeros_like(param).cuda(param.device)
        )
//  [python]
def test_averaged_model_state_dict(self) -> None:
        dnn = torch.nn.Sequential(
            torch.nn.Conv2d(1, 5, kernel_size=3), torch.nn.Linear(5, 10)
        )
        averaged_dnn = AveragedModel(dnn)
        averaged_dnn2 = AveragedModel(dnn)
        n_updates = 10
        for _ in range(n_updates):
            for p in dnn.parameters():
                p.detach().add_(torch.randn_like(p))
            averaged_dnn.update_parameters(dnn)
        averaged_dnn2.load_state_dict(averaged_dnn.state_dict())
        for p_swa, p_swa2 in zip(averaged_dnn.parameters(), averaged_dnn2.parameters()):
            self.assertEqual(p_swa, p_swa2)
        self.assertTrue(averaged_dnn.n_averaged == averaged_dnn2.n_averaged)
//  [python]
def test_averaged_model_state_dict(self) -> None:
        dnn = torch.nn.Sequential(
            torch.nn.Conv2d(1, 5, kernel_size=3), torch.nn.Linear(5, 10)
        )
        averaged_dnn = AveragedModel(dnn)
        averaged_dnn2 = AveragedModel(dnn)
        n_updates = 10
        for _ in range(n_updates):
            for p in dnn.parameters():
                p.detach().add_(torch.randn_like(p))
            averaged_dnn.update_parameters(dnn)
        averaged_dnn2.load_state_dict(averaged_dnn.state_dict())
        for p_swa, p_swa2 in zip(averaged_dnn.parameters(), averaged_dnn2.parameters()):
            torch.testing.assert_close(p_swa, p_swa2, check_device=False)
        self.assertTrue(averaged_dnn.n_averaged == averaged_dnn2.n_averaged)
//  [python]
def main():
    session = brt.Session(alloc_func=caching_allocator_alloc, free_func=caching_allocator_delete)
    ait_dir = os.path.join(os.path.dirname(__file__), "../../test/test_files/AITOp")
    model_path = os.path.join(ait_dir, "bmm_permute_entry.mlir")
    session.load(model_path)
    req = session.new_request_context(torch.cuda.current_stream()._as_parameter_.value)

    inputs = []
    outputs = []

    for offset in session.get_input_arg_offsets():
        data = np.random.random(size=session.get_static_shape(offset))
        inputs.append(torch.tensor(data, dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, inputs[-1].data_ptr())

    for offset in session.get_output_arg_offsets():
        outputs.append(torch.empty(session.get_static_shape(offset), dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, outputs[-1].data_ptr())

    req.finish_io_binding()
    req.run()
    req.sync()

    b = inputs[0].size(dim=0)
    m = inputs[0].size(dim=1)
    N = inputs[1].size(dim=2)
    d1 = outputs[0].size(dim=2)

    Y_l = torch.bmm(inputs[0], inputs[1])
    Y_r = Y_l.reshape(b // d1, d1, m, N)
    Y_pt = torch.permute(Y_r, [0, 2, 1, 3])

    torch.testing.assert_allclose(Y_pt, outputs[0], atol=1e-3, rtol=1e-3)
    print("bmm_permute numerical test pass")
//  [python]
def main():
    session = brt.Session(alloc_func=caching_allocator_alloc, free_func=caching_allocator_delete)
    ait_dir = os.path.join(os.path.dirname(__file__), "../../test/test_files/AITOp")
    model_path = os.path.join(ait_dir, "bmm_permute_entry.mlir")
    session.load(model_path)
    req = session.new_request_context(torch.cuda.current_stream()._as_parameter_.value)

    inputs = []
    outputs = []

    for offset in session.get_input_arg_offsets():
        data = np.random.random(size=session.get_static_shape(offset))
        inputs.append(torch.tensor(data, dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, inputs[-1].data_ptr())

    for offset in session.get_output_arg_offsets():
        outputs.append(torch.empty(session.get_static_shape(offset), dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, outputs[-1].data_ptr())

    req.finish_io_binding()
    req.run()
    req.sync()

    b = inputs[0].size(dim=0)
    m = inputs[0].size(dim=1)
    N = inputs[1].size(dim=2)
    d1 = outputs[0].size(dim=2)

    Y_l = torch.bmm(inputs[0], inputs[1])
    Y_r = Y_l.reshape(b // d1, d1, m, N)
    Y_pt = torch.permute(Y_r, [0, 2, 1, 3])

    torch.testing.assert_close(Y_pt, outputs[0], atol=1e-3, rtol=1e-3)
    print("bmm_permute numerical test pass")
//  [python]
def test_rosen():
    """Test Rosenbrock problem with constraints."""

    x0 = torch.tensor([1., 8.])
    res = minimize(
        rosen, x0,
        method='l-bfgs',
        options=dict(line_search='strong-wolfe'),
        max_iter=50,
        disp=0
    )


    # Test inactive constraints

    res_constrained_sum = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.sum(), ub=10.),
        max_iter=50,
        disp=0
    )
    t.assert_close(
        res.x, res_constrained_sum.x, rtol=1e-2, atol=1e-2)

    res_constrained_norm = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.square().sum(), ub=10.),
        max_iter=50,
        disp=0
    )
    t.assert_close(
        res.x, res_constrained_norm.x, rtol=1e-2, atol=1e-2)


    # Test active constraints

    res_constrained_sum = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.sum(), ub=1.),
        max_iter=50,
        disp=0
    )
    assert res_constrained_sum.x.sum() <= 1.
    res_constrained_norm = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.square().sum(), ub=1.),
        max_iter=50,
        disp=0
    )
    assert res_constrained_norm.x.square().sum() <= 1.
//  [python]
def test_rosen():
    """Test Rosenbrock problem with constraints."""

    x0 = torch.tensor([1., 8.])
    res = minimize(
        rosen, x0,
        method='l-bfgs',
        options=dict(line_search='strong-wolfe'),
        max_iter=50,
        disp=0
    )


    # Test inactive constraints

    res_constrained_sum = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.sum(), ub=10.),
        max_iter=50,
        disp=0
    )
    torch.testing.assert_close(
        res.x, res_constrained_sum.x, rtol=1e-2, atol=1e-2)

    res_constrained_norm = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.square().sum(), ub=10.),
        max_iter=50,
        disp=0
    )
    torch.testing.assert_close(
        res.x, res_constrained_norm.x, rtol=1e-2, atol=1e-2)


    # Test active constraints

    res_constrained_sum = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.sum(), ub=1.),
        max_iter=50,
        disp=0
    )
    assert res_constrained_sum.x.sum() <= 1.
    res_constrained_norm = minimize_constr(
        rosen, x0,
        constr=dict(fun=lambda x: x.square().sum(), ub=1.),
        max_iter=50,
        disp=0
    )
    assert res_constrained_norm.x.square().sum() <= 1.
//  [python]
def _test_averaged_model(
        self, net_device: torch.device, swa_device: torch.device, ema: bool
    ) -> None:
        dnn = torch.nn.Sequential(
            torch.nn.Conv2d(1, 5, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
            torch.nn.BatchNorm2d(5, momentum=0.3),
            torch.nn.Conv2d(5, 2, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 10),
        ).to(net_device)

        averaged_params, averaged_dnn = self._run_averaged_steps(dnn, swa_device, ema)

        for p_avg, p_swa in zip(averaged_params, averaged_dnn.parameters()):
            torch.allclose(p_avg, p_swa, check_device=False)
            # Check that AveragedModel is on the correct device
            self.assertTrue(p_swa.device == swa_device)
            self.assertTrue(p_avg.device == net_device)
        self.assertTrue(averaged_dnn.n_averaged.device == swa_device)
//  [python]
def _test_averaged_model(
        self, net_device: torch.device, swa_device: torch.device, ema: bool
    ) -> None:
        dnn = torch.nn.Sequential(
            torch.nn.Conv2d(1, 5, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
            torch.nn.BatchNorm2d(5, momentum=0.3),
            torch.nn.Conv2d(5, 2, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 10),
        ).to(net_device)

        averaged_params, averaged_dnn = self._run_averaged_steps(dnn, swa_device, ema)

        for p_avg, p_swa in zip(averaged_params, averaged_dnn.parameters()):
            torch.testing.assert_close(p_avg, p_swa, check_device=False)
            # Check that AveragedModel is on the correct device
            self.assertTrue(p_swa.device == swa_device)
            self.assertTrue(p_avg.device == net_device)
        self.assertTrue(averaged_dnn.n_averaged.device == swa_device)
//  [python]
def test_lit_ema(self) -> None:
        model = torch.nn.Linear(10, 10)
        ema_decay = 0.999
        averaged_model = AveragedModel(
            model,
            averaging_method="ema",
            ema_decay=ema_decay,
            use_lit=True,
        )

        averaged_params = [torch.zeros_like(param) for param in model.parameters()]

        n_updates = 10
        for i in range(n_updates):
            decay = min(ema_decay, (1 + i + 1) / (10 + i + 1))

            updated_averaged_params = []
            for p, p_avg in zip(model.parameters(), averaged_params):
                p.detach().add_(torch.randn_like(p))
                if i == 0:
                    updated_averaged_params.append(p.clone())
                else:
                    updated_averaged_params.append(
                        (p_avg * decay + p * (1 - decay)).clone()
                    )
            averaged_model.update_parameters(model)
            averaged_params = updated_averaged_params

        for p_avg, p_swa in zip(averaged_params, averaged_model.parameters()):
            self.assertEqual(p_avg, p_swa)
//  [python]
def test_lit_ema(self) -> None:
        model = torch.nn.Linear(10, 10)
        ema_decay = 0.999
        averaged_model = AveragedModel(
            model,
            averaging_method="ema",
            ema_decay=ema_decay,
            use_lit=True,
        )

        averaged_params = [torch.zeros_like(param) for param in model.parameters()]

        n_updates = 10
        for i in range(n_updates):
            decay = min(ema_decay, (1 + i + 1) / (10 + i + 1))

            updated_averaged_params = []
            for p, p_avg in zip(model.parameters(), averaged_params):
                p.detach().add_(torch.randn_like(p))
                if i == 0:
                    updated_averaged_params.append(p.clone())
                else:
                    updated_averaged_params.append(
                        (p_avg * decay + p * (1 - decay)).clone()
                    )
            averaged_model.update_parameters(model)
            averaged_params = updated_averaged_params

        for p_avg, p_swa in zip(averaged_params, averaged_model.parameters()):
            torch.testing.assert_close(p_avg, p_swa, check_device=False)
//  [python]
def test_to_tensor(self, test_scripting, configs):
        """test tensorization on both single sequence and batch of sequence"""
        inputs, padding_value, expected_list = configs
        func = functional.to_tensor
        if test_scripting:
            func = torch.jit.script(func)

        actual = func(inputs, padding_value=padding_value)
        expected = torch.tensor(expected_list, dtype=torch.long)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_to_tensor(self, test_scripting, configs):
        """test tensorization on both single sequence and batch of sequence"""
        inputs, padding_value, expected_list = configs
        func = functional.to_tensor
        if test_scripting:
            func = torch.jit.script(func)

        actual = func(inputs, padding_value=padding_value)
        expected = torch.tensor(expected_list, dtype=torch.long)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_get_prompt_logprobs(
    hf_runner,
    aphrodite_runner,
    model,
    dtype,
    example_prompts,
):
    max_tokens = 5
    hf_model = hf_runner(model, dtype=dtype)
    hf_logprobs = hf_model.generate_greedy_logprobs(
        example_prompts,
        max_tokens=max_tokens,
    )
    del hf_model

    aphrodite_model = aphrodite_runner(model, dtype=dtype)
    aphrodite_sampling_params = SamplingParams(max_tokens=max_tokens,
                                               logprobs=5,
                                               prompt_logprobs=5,
                                               temperature=0.0)
    aphrodite_results = aphrodite_model.model.generate(
        example_prompts, sampling_params=aphrodite_sampling_params)

    # Test whether logprobs are included in the results.
    for result in aphrodite_results:
        assert result.prompt_logprobs is not None
        assert result.outputs[0].logprobs is not None

    # Test whether prompt logprobs are consistent with HF
    for aphrodite_result, hf_logprob in zip(aphrodite_results, hf_logprobs):
        # Check prompt logprobs
        aphrodite_prompt_logprobs = aphrodite_result.prompt_logprobs[1:]
        for i, aphrodite_prompt_logprob_dict in enumerate(
                aphrodite_prompt_logprobs):
            for token_id, logprob in aphrodite_prompt_logprob_dict.items():
                torch.testing.assert_allclose(logprob,
                                           hf_logprob[0][i][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
        aphrodite_sample_logprobs = aphrodite_result.outputs[0].logprobs
        for i, aphrodite_sample_logprob_dict in enumerate(
                aphrodite_sample_logprobs):
            for token_id, logprob in aphrodite_sample_logprob_dict.items():
                torch.testing.assert_allclose(logprob,
                                           hf_logprob[i][-1][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
//  [python]
def test_get_prompt_logprobs(
    hf_runner,
    aphrodite_runner,
    model,
    dtype,
    example_prompts,
):
    max_tokens = 5
    hf_model = hf_runner(model, dtype=dtype)
    hf_logprobs = hf_model.generate_greedy_logprobs(
        example_prompts,
        max_tokens=max_tokens,
    )
    del hf_model

    aphrodite_model = aphrodite_runner(model, dtype=dtype)
    aphrodite_sampling_params = SamplingParams(max_tokens=max_tokens,
                                               logprobs=5,
                                               prompt_logprobs=5,
                                               temperature=0.0)
    aphrodite_results = aphrodite_model.model.generate(
        example_prompts, sampling_params=aphrodite_sampling_params)

    # Test whether logprobs are included in the results.
    for result in aphrodite_results:
        assert result.prompt_logprobs is not None
        assert result.outputs[0].logprobs is not None

    # Test whether prompt logprobs are consistent with HF
    for aphrodite_result, hf_logprob in zip(aphrodite_results, hf_logprobs):
        # Check prompt logprobs
        aphrodite_prompt_logprobs = aphrodite_result.prompt_logprobs[1:]
        for i, aphrodite_prompt_logprob_dict in enumerate(
                aphrodite_prompt_logprobs):
            for token_id, logprob in aphrodite_prompt_logprob_dict.items():
                torch.testing.assert_close(logprob,
                                           hf_logprob[0][i][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
        aphrodite_sample_logprobs = aphrodite_result.outputs[0].logprobs
        for i, aphrodite_sample_logprob_dict in enumerate(
                aphrodite_sample_logprobs):
            for token_id, logprob in aphrodite_sample_logprob_dict.items():
                torch.testing.assert_close(logprob,
                                           hf_logprob[i][-1][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
//  [python]
def test_forward(self, inputs, lora_linear, out_dim) -> None:
        expected = torch.tensor(EXPECTED_VAL)
        actual = lora_linear(inputs)
        assert actual.shape == (BSZ, SEQ_LEN, out_dim)
        torch.testing.assert_allclose(actual.mean(), expected, atol=1e-4, rtol=1e-6)
//  [python]
def test_forward(self, inputs, lora_linear, out_dim) -> None:
        expected = torch.tensor(EXPECTED_VAL)
        actual = lora_linear(inputs)
        assert actual.shape == (BSZ, SEQ_LEN, out_dim)
        torch.testing.assert_close(actual.mean(), expected, atol=1e-4, rtol=1e-6)
//  [python]
def test_rng_decorator(device):
    with torch.device(device):
        torch.manual_seed(10)
        s0a = torch.randn(3)
        with _rng_decorator(0):
            torch.randn(3)
        s0b = torch.randn(3)
        torch.manual_seed(10)
        s1a = torch.randn(3)
        s1b = torch.randn(3)
        t.assert_close(s0a, s1a)
        t.assert_close(s0b, s1b)
//  [python]
def test_rng_decorator(device):
    with torch.device(device):
        torch.manual_seed(10)
        s0a = torch.randn(3)
        with _rng_decorator(0):
            torch.randn(3)
        s0b = torch.randn(3)
        torch.manual_seed(10)
        s1a = torch.randn(3)
        s1b = torch.randn(3)
        torch.testing.assert_close(s0a, s1a)
        torch.testing.assert_close(s0b, s1b)
//  [python]
def test_qlora_parity(self, dtype):
        with utils.set_default_dtype(dtype):
            qlora_linear = LoRALinear(
                in_dim=512,
                out_dim=512,
                rank=RANK,
                alpha=ALPHA,
                use_bias=False,
                quantize_base=True,
            )
            lora_linear = LoRALinear(
                in_dim=512,
                out_dim=512,
                rank=RANK,
                alpha=ALPHA,
                use_bias=False,
                quantize_base=False,
            )

        # set weight of lora_linear to unquantized weight of qlora_linear and check
        # parity.
        lora_linear.weight.data = qlora_linear.weight.to(dtype)

        # Ensure forward passes are the same. This is because LoRALinear should use a special
        # quantized linear operator that runs compute in higher prec (but only saves the 4 bit quantized tensor)
        # for autograd.
        inputs = torch.randn(BSZ, SEQ_LEN, 512, dtype=dtype)
        lora_linear_out = lora_linear(inputs)
        qlora_linear_out = qlora_linear(inputs)
        torch.testing.assert_allclose(lora_linear_out, qlora_linear_out)
//  [python]
def test_qlora_parity(self, dtype):
        with utils.set_default_dtype(dtype):
            qlora_linear = LoRALinear(
                in_dim=512,
                out_dim=512,
                rank=RANK,
                alpha=ALPHA,
                use_bias=False,
                quantize_base=True,
            )
            lora_linear = LoRALinear(
                in_dim=512,
                out_dim=512,
                rank=RANK,
                alpha=ALPHA,
                use_bias=False,
                quantize_base=False,
            )

        # set weight of lora_linear to unquantized weight of qlora_linear and check
        # parity.
        lora_linear.weight.data = qlora_linear.weight.to(dtype)

        # Ensure forward passes are the same. This is because LoRALinear should use a special
        # quantized linear operator that runs compute in higher prec (but only saves the 4 bit quantized tensor)
        # for autograd.
        inputs = torch.randn(BSZ, SEQ_LEN, 512, dtype=dtype)
        lora_linear_out = lora_linear(inputs)
        qlora_linear_out = qlora_linear(inputs)
        torch.testing.assert_close(lora_linear_out, qlora_linear_out)
//  [python]
def test_causal_model_logits(self):
        """comparing logits outputs of whole inner model"""
        # Input 0: one row per sentence; Input 1: same data, but stacked into a single row with custom attention
        input_0, position_ids_0, input_1, mask_1, position_ids_1 = self.get_test_data()

        logits_0 = self.model.forward(input_0, position_ids=position_ids_0).logits
        logits_1 = self.model.forward(input_1, attention_mask=mask_1.bool(), position_ids=position_ids_1).logits

        logits_0_last_tokens = logits_0[:, -1, :]  # last tokens in each batch line
        logits_1_last_tokens = logits_1[0, -3:, :]  # last three tokens
        torch.testing.assert_allclose(logits_0_last_tokens, logits_1_last_tokens)
//  [python]
def test_causal_model_logits(self):
        """comparing logits outputs of whole inner model"""
        # Input 0: one row per sentence; Input 1: same data, but stacked into a single row with custom attention
        input_0, position_ids_0, input_1, mask_1, position_ids_1 = self.get_test_data()

        logits_0 = self.model.forward(input_0, position_ids=position_ids_0).logits
        logits_1 = self.model.forward(input_1, attention_mask=mask_1.bool(), position_ids=position_ids_1).logits

        logits_0_last_tokens = logits_0[:, -1, :]  # last tokens in each batch line
        logits_1_last_tokens = logits_1[0, -3:, :]  # last three tokens
        torch.testing.assert_close(logits_0_last_tokens, logits_1_last_tokens)
//  [python]
def test_binary_add_tensors():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_binary_add_tensors():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    assert_close(actual, expected)
//  [python]
def test_trainer(recipe_name):
    with open('yamls/resnet50.yaml') as f:
        base_config = OmegaConf.load(f)

    with open('tests/smoketest_config.yaml') as f:
        smoke_config = OmegaConf.load(f)
    config = OmegaConf.merge(base_config, smoke_config)
    config.recipe_name = recipe_name

    with SynthClassificationDirectory() as tmp_datadir:
        print(tmp_datadir)
        config.train_dataset.path = tmp_datadir
        config.eval_dataset.path = tmp_datadir
        # Also save checkpoints in the temporary directory
        config.save_folder = tmp_datadir

        # Train
        trainer1 = main(config)
        model1 = trainer1.state.model.module

        # TODO avoid tests taking a long time to exit without this
        trainer1.state.dataloader._iterator._shutdown_workers()  # type: ignore

        # Check that the checkpoint was saved
        chkpt_path = os.path.join(tmp_datadir, 'ep0-ba1-rank0.pt')
        assert os.path.isfile(chkpt_path)

        # Check that the checkpoint was loaded by comparing model weights
        config.load_path = chkpt_path
        config.is_train = False
        config.seed += 10  # change seed
        trainer2 = main(config)
        model2 = trainer2.state.model.module

        for param1, param2 in zip(model1.parameters(), model2.parameters()):
            torch.testing.assert_allclose(param1, param2)
//  [python]
def test_trainer(recipe_name):
    with open('yamls/resnet50.yaml') as f:
        base_config = OmegaConf.load(f)

    with open('tests/smoketest_config.yaml') as f:
        smoke_config = OmegaConf.load(f)
    config = OmegaConf.merge(base_config, smoke_config)
    config.recipe_name = recipe_name

    with SynthClassificationDirectory() as tmp_datadir:
        print(tmp_datadir)
        config.train_dataset.path = tmp_datadir
        config.eval_dataset.path = tmp_datadir
        # Also save checkpoints in the temporary directory
        config.save_folder = tmp_datadir

        # Train
        trainer1 = main(config)
        model1 = trainer1.state.model.module

        # TODO avoid tests taking a long time to exit without this
        trainer1.state.dataloader._iterator._shutdown_workers()  # type: ignore

        # Check that the checkpoint was saved
        chkpt_path = os.path.join(tmp_datadir, 'ep0-ba1-rank0.pt')
        assert os.path.isfile(chkpt_path)

        # Check that the checkpoint was loaded by comparing model weights
        config.load_path = chkpt_path
        config.is_train = False
        config.seed += 10  # change seed
        trainer2 = main(config)
        model2 = trainer2.state.model.module

        for param1, param2 in zip(model1.parameters(), model2.parameters()):
            torch.testing.assert_close(param1, param2)
//  [python]
def test_attention(self):
        """comparing outputs of attention layer"""
        # Input 0: one row per sentence; Input 1: same data, but stacked into a single row with custom attention
        input_0, position_ids_0, input_1, mask_1, position_ids_1 = self.get_test_data()
        causal_mask_1 = (1 - mask_1).to(self.model_dtype) * torch.finfo(self.model_dtype).min

        hid_0 = self.model.model.embed_tokens(input_0)
        outs_0 = self.model.model.layers[0].self_attn.forward(hid_0, position_ids=position_ids_0)[0]
        # outs_0.shape == torch.Size([3, 4, 768])

        hid_1 = self.model.model.embed_tokens(input_1)
        outs_1 = self.model.model.layers[0].self_attn.forward(
            hid_1, attention_mask=causal_mask_1, position_ids=position_ids_1
        )[0]
        # outs_1.shape == torch.Size([1, 6, 768])

        outs_0_last_tokens = outs_0[:, -1, :]  # last tokens in each batch line
        outs_1_last_tokens = outs_1[0, -3:, :]  # last three tokens
        self.assertEqual(outs_0_last_tokens, outs_1_last_tokens)
//  [python]
def test_attention(self):
        """comparing outputs of attention layer"""
        # Input 0: one row per sentence; Input 1: same data, but stacked into a single row with custom attention
        input_0, position_ids_0, input_1, mask_1, position_ids_1 = self.get_test_data()
        causal_mask_1 = (1 - mask_1).to(self.model_dtype) * torch.finfo(self.model_dtype).min

        hid_0 = self.model.model.embed_tokens(input_0)
        outs_0 = self.model.model.layers[0].self_attn.forward(hid_0, position_ids=position_ids_0)[0]
        # outs_0.shape == torch.Size([3, 4, 768])

        hid_1 = self.model.model.embed_tokens(input_1)
        outs_1 = self.model.model.layers[0].self_attn.forward(
            hid_1, attention_mask=causal_mask_1, position_ids=position_ids_1
        )[0]
        # outs_1.shape == torch.Size([1, 6, 768])

        outs_0_last_tokens = outs_0[:, -1, :]  # last tokens in each batch line
        outs_1_last_tokens = outs_1[0, -3:, :]  # last three tokens
        torch.testing.assert_close(outs_0_last_tokens, outs_1_last_tokens)
//  [python]
def test_tower_dropout(self):
        blocks = []
        in_planes = 8
        out_planes = 8
        for _ in range(2):
            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))

        block = jasper.ParallelBlock(blocks, aggregation_mode='dropout', block_dropout_prob=1.0)
        x = torch.randn(1, in_planes, 140)
        xlen = torch.tensor([131])
        y, _ = block(([x], xlen))

        # Tower dropout is 1.0, meaning that all towers have to be dropped, so only residual remains.
        self.assertEqual(y[0], x)
//  [python]
def test_tower_dropout(self):
        blocks = []
        in_planes = 8
        out_planes = 8
        for _ in range(2):
            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))

        block = jasper.ParallelBlock(blocks, aggregation_mode='dropout', block_dropout_prob=1.0)
        x = torch.randn(1, in_planes, 140)
        xlen = torch.tensor([131])
        y, _ = block(([x], xlen))

        # Tower dropout is 1.0, meaning that all towers have to be dropped, so only residual remains.
        torch.testing.assert_close(y[0], x)
//  [python]
def test_torch_add_tensors():
    def foo(a, b):
        return torch.add(a, b)

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_torch_add_tensors():
    def foo(a, b):
        return torch.add(a, b)

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    assert_close(actual, expected)
//  [python]
def test_causal_model_logits(self):
        """comparing logits outputs of whole inner model"""
        # Input 0: one row per sentence; Input 1: same data, but stacked into a single row with custom attention
        input_0, position_ids_0, input_1, mask_1, position_ids_1 = self.get_test_data()

        logits_0 = self.model.forward(input_0, position_ids=position_ids_0).logits
        logits_1 = self.model.forward(input_1, attention_mask=mask_1.bool(), position_ids=position_ids_1).logits

        logits_0_last_tokens = logits_0[:, -1, :]  # last tokens in each batch line
        logits_1_last_tokens = logits_1[0, -3:, :]  # last three tokens

        indices_0 = logits_0_last_tokens.sort(descending=True).indices
        indices_1 = logits_1_last_tokens.sort(descending=True).indices

        # checking logits, but note relaxed tolerances for FP16
        self.assertTrue(torch.allclose(logits_0_last_tokens, logits_1_last_tokens, atol=0.02, rtol=0.001))

        # checking tokens order for the top tokens
        for token_ids_0, token_ids_1 in zip(indices_0, indices_1):
            self.assertTrue(torch.equal(token_ids_0[:128], token_ids_1[:128]))
//  [python]
def test_causal_model_logits(self):
        """comparing logits outputs of whole inner model"""
        # Input 0: one row per sentence; Input 1: same data, but stacked into a single row with custom attention
        input_0, position_ids_0, input_1, mask_1, position_ids_1 = self.get_test_data()

        logits_0 = self.model.forward(input_0, position_ids=position_ids_0).logits
        logits_1 = self.model.forward(input_1, attention_mask=mask_1.bool(), position_ids=position_ids_1).logits

        logits_0_last_tokens = logits_0[:, -1, :]  # last tokens in each batch line
        logits_1_last_tokens = logits_1[0, -3:, :]  # last three tokens

        indices_0 = logits_0_last_tokens.sort(descending=True).indices
        indices_1 = logits_1_last_tokens.sort(descending=True).indices

        # checking logits, but note relaxed tolerances for FP16
        torch.testing.assert_close(logits_0_last_tokens, logits_1_last_tokens, atol=0.02, rtol=0.001)

        # checking tokens order for the top tokens
        for token_ids_0, token_ids_1 in zip(indices_0, indices_1):
            self.assertTrue(torch.equal(token_ids_0[:128], token_ids_1[:128]))
//  [python]
def test_brax_consistency(self, envname, batch_size, requires_grad):
        import jax
        import jax.numpy as jnp
        from torchrl.envs.libs.jax_utils import (
            _ndarray_to_tensor,
            _tensor_to_ndarray,
            _tree_flatten,
        )

        env = BraxEnv(envname, batch_size=batch_size, requires_grad=requires_grad)
        env.set_seed(1)
        rollout = env.rollout(10)

        env.set_seed(1)
        key = env._key
        base_env = env._env
        key, *keys = jax.random.split(key, int(np.prod(batch_size) + 1))
        state = jax.vmap(base_env.reset)(jnp.stack(keys))
        for i in range(rollout.shape[-1]):
            action = rollout[..., i]["action"]
            action = _tensor_to_ndarray(action.clone())
            action = _tree_flatten(action, env.batch_size)
            state = jax.vmap(base_env.step)(state, action)
            t1 = rollout[..., i][("next", "observation")]
            t2 = _ndarray_to_tensor(state.obs).view_as(t1)
            torch.testing.assert_allclose(t1, t2)
//  [python]
def test_brax_consistency(self, envname, batch_size, requires_grad):
        import jax
        import jax.numpy as jnp
        from torchrl.envs.libs.jax_utils import (
            _ndarray_to_tensor,
            _tensor_to_ndarray,
            _tree_flatten,
        )

        env = BraxEnv(envname, batch_size=batch_size, requires_grad=requires_grad)
        env.set_seed(1)
        rollout = env.rollout(10)

        env.set_seed(1)
        key = env._key
        base_env = env._env
        key, *keys = jax.random.split(key, int(np.prod(batch_size) + 1))
        state = jax.vmap(base_env.reset)(jnp.stack(keys))
        for i in range(rollout.shape[-1]):
            action = rollout[..., i]["action"]
            action = _tensor_to_ndarray(action.clone())
            action = _tree_flatten(action, env.batch_size)
            state = jax.vmap(base_env.step)(state, action)
            t1 = rollout[..., i][("next", "observation")]
            t2 = _ndarray_to_tensor(state.obs).view_as(t1)
            torch.testing.assert_close(t1, t2)
//  [python]
def test_torch_add_tensors_closure():
    def foo(a, b):
        c = a + b

        def bar():
            return torch.add(c, 1)

        return bar()

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    t.assert_close(actual, expected)
//  [python]
def test_torch_add_tensors_closure():
    def foo(a, b):
        c = a + b

        def bar():
            return torch.add(c, 1)

        return bar()

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    assert_close(actual, expected)
//  [python]
def test_torch_add_tensors_closure_external():
    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    def bar(b):
        return torch.add(a, b)

    def foo():
        bar(b)

    jbar = thunder.jit(bar)
    actual = jbar(b)
    expected = bar(b)
    torch.testing.assert_allclose(actual, expected)

    jfoo = thunder.jit(foo)
    actual = jfoo()
    expected = foo()
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_torch_add_tensors_closure_external():
    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    def bar(b):
        return torch.add(a, b)

    def foo():
        bar(b)

    jbar = thunder.jit(bar)
    actual = jbar(b)
    expected = bar(b)
    assert_close(actual, expected)

    jfoo = thunder.jit(foo)
    actual = jfoo()
    expected = foo()
    assert_close(actual, expected)
//  [python]
def test_intermediate_torch_operations():
    def foo(a, b):
        c = a + b
        d = torch.sub(c, b)
        e = torch.mul(d, a)
        f = torch.matmul(e, c)
        g = [e, f]
        return torch.cat(g)

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_intermediate_torch_operations():
    def foo(a, b):
        c = a + b
        d = torch.sub(c, b)
        e = torch.mul(d, a)
        f = torch.matmul(e, c)
        g = [e, f]
        return torch.cat(g)

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    actual = jfoo(a, b)
    expected = foo(a, b)

    assert_close(actual, expected)
//  [python]
def test_cache_always_trace():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo, cache=CACHE_OPTIONS.NO_CACHING)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    expected = foo(a, b)
    actual = jfoo(a, b)
    actual = jfoo(a, b)
    actual = jfoo(a, b)
    actual = jfoo(a, b)
    t.assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 4
    assert thunder.cache_hits(jfoo) == 0
//  [python]
def test_cache_always_trace():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo, cache=CACHE_OPTIONS.NO_CACHING)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    expected = foo(a, b)
    actual = jfoo(a, b)
    actual = jfoo(a, b)
    actual = jfoo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 4
    assert thunder.cache_hits(jfoo) == 0
//  [python]
def test_nn_parameter():
    a = torch.nn.Parameter(torch.randn(2, 3))
    b = torch.tensor(2)

    def fn(a):
        return b * a

    jfn = thunder.jit(fn)

    expected = fn(a)
    actual = jfn(a)
    t.assert_close(expected, actual)
//  [python]
def test_nn_parameter():
    a = torch.nn.Parameter(torch.randn(2, 3))
    b = torch.tensor(2)

    def fn(a):
        return b * a

    jfn = thunder.jit(fn)

    expected = fn(a)
    actual = jfn(a)
    assert_close(expected, actual)
//  [python]
def test_cache_equality_constraint():
    x, y = torch.randn(2, 2)

    def fn(b):
        if b:
            return x
        else:
            return y

    jfn = thunder.jit(fn)

    t.assert_close(fn(True), jfn(True))
    t.assert_close(fn(False), jfn(False))

    assert thunder.cache_misses(jfn) == 2
    assert thunder.cache_hits(jfn) == 0

    t.assert_close(fn(True), jfn(True))
    t.assert_close(fn(False), jfn(False))

    assert thunder.cache_misses(jfn) == 2
    assert thunder.cache_hits(jfn) == 2
//  [python]
def test_cache_equality_constraint():
    x, y = torch.randn(2, 2)

    def fn(b):
        if b:
            return x
        else:
            return y

    jfn = thunder.jit(fn)

    assert_close(fn(True), jfn(True))
    assert_close(fn(False), jfn(False))

    assert thunder.cache_misses(jfn) == 2
    assert thunder.cache_hits(jfn) == 0

    assert_close(fn(True), jfn(True))
    assert_close(fn(False), jfn(False))

    assert thunder.cache_misses(jfn) == 2
    assert thunder.cache_hits(jfn) == 2
//  [python]
def test_add_numbers():
    def foo(a, b):
        return torch.add(a, b)

    jfoo = thunder.jit(foo)

    # TODO Add test for bool
    # see issue "Binary addition on booleans should promote to an integer"
    cases = (
        (2, 3),
        (2.1, 3.4),
        (complex(1, 1), complex(-1, 2)),
    )

    for a, b in cases:
        actual = jfoo(a, b)
        expected = a + b

        t.assert_close(actual, expected)
//  [python]
def test_add_numbers():
    def foo(a, b):
        return torch.add(a, b)

    jfoo = thunder.jit(foo)

    # TODO Add test for bool
    # see issue "Binary addition on booleans should promote to an integer"
    cases = (
        (2, 3),
        (2.1, 3.4),
        (complex(1, 1), complex(-1, 2)),
    )

    for a, b in cases:
        actual = jfoo(a, b)
        expected = a + b

        assert_close(actual, expected)
//  [python]
def test_binary_add_tensor_number():
    # Tests using torch.add
    def foo(a):
        return torch.add(a, 3)

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")

    actual = jfoo(a)
    expected = foo(a)

    t.assert_close(actual, expected)

    # Tests using addition operator
    def foo(a):
        return a + 4

    jfoo = thunder.jit(foo)

    actual = jfoo(a)
    expected = foo(a)

    t.assert_close(actual, expected)
//  [python]
def test_binary_add_tensor_number():
    # Tests using torch.add
    def foo(a):
        return torch.add(a, 3)

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")

    actual = jfoo(a)
    expected = foo(a)

    assert_close(actual, expected)

    # Tests using addition operator
    def foo(a):
        return a + 4

    jfoo = thunder.jit(foo)

    actual = jfoo(a)
    expected = foo(a)

    assert_close(actual, expected)
//  [python]
def test_binary_add_numbers():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo)

    # TODO Add test for bool
    # see issue "Binary addition on booleans should promote to an integer"
    cases = (
        (2, 3),
        (2.1, 3.4),
        (complex(1, 1), complex(-1, 2)),
    )

    for a, b in cases:
        actual = jfoo(a, b)
        expected = foo(a, b)

        t.assert_close(actual, expected)
//  [python]
def test_binary_add_numbers():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo)

    # TODO Add test for bool
    # see issue "Binary addition on booleans should promote to an integer"
    cases = (
        (2, 3),
        (2.1, 3.4),
        (complex(1, 1), complex(-1, 2)),
    )

    for a, b in cases:
        actual = jfoo(a, b)
        expected = foo(a, b)

        assert_close(actual, expected)
//  [python]
def test_lookaside_bool():
    def foo(a, b, i):
        if bool(i):
            return a + b
        return a - b

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    expected = foo(a, b, 0)
    actual = jfoo(a, b, 0)
    t.assert_close(expected, actual)

    expected = foo(a, b, 1)
    actual = jfoo(a, b, 1)
    t.assert_close(expected, actual)