jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_tl_range(device):
    if is_hip():
        pytest.skip("test_tl_range is not supported in HIP")
    if is_interpreter():
        pytest.skip("FIXME: RuntimeError: \"addmm_impl_cpu_\" not implemented for 'Half'")
    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 testConvolutionBackward(self):
        test_device = get_test_device()
        inp = torch.rand(1, 3, 128, 128, device=test_device, requires_grad=True)
        inp_copy = clone_move(inp)
        grad = torch.rand(1, 32, 121, 121, device=test_device)  # no requires_grad
        grad_copy = clone_move(grad)
        weight = torch.rand(32, 3, 8, 8, device=test_device, requires_grad=True)
        weight_copy = clone_move(weight)
        bias = torch.rand(32, device=test_device, requires_grad=True)
        bias_copy = clone_move(bias)

        # run eager
        conv_out = torch.nn.functional.conv2d(inp, weight, bias)
        (inp_grad, weight_grad, bias_grad) = torch.autograd.grad(
            [conv_out], [inp, weight, bias], [grad]
        )

        # run lazy
        conv_copy_out = torch.nn.functional.conv2d(inp_copy, weight_copy, bias_copy)
        (inp_copy_grad, weight_copy_grad, bias_copy_grad) = torch.autograd.grad(
            [conv_copy_out], [inp_copy, weight_copy, bias_copy], [grad_copy]
        )

        # check numerics
        t.assertEqual(bias_copy_grad.cpu(), bias_grad.cpu())

        t.assertEqual(weight_copy_grad.cpu(), weight_grad.cpu())
        t.assertEqual(inp_copy_grad.cpu(), inp_grad.cpu())
//  [python]
def testConvolutionBackward(self):
        test_device = get_test_device()
        inp = torch.rand(1, 3, 128, 128, device=test_device, requires_grad=True)
        inp_copy = clone_move(inp)
        grad = torch.rand(1, 32, 121, 121, device=test_device)  # no requires_grad
        grad_copy = clone_move(grad)
        weight = torch.rand(32, 3, 8, 8, device=test_device, requires_grad=True)
        weight_copy = clone_move(weight)
        bias = torch.rand(32, device=test_device, requires_grad=True)
        bias_copy = clone_move(bias)

        # run eager
        conv_out = torch.nn.functional.conv2d(inp, weight, bias)
        (inp_grad, weight_grad, bias_grad) = torch.autograd.grad(
            [conv_out], [inp, weight, bias], [grad]
        )

        # run lazy
        conv_copy_out = torch.nn.functional.conv2d(inp_copy, weight_copy, bias_copy)
        (inp_copy_grad, weight_copy_grad, bias_copy_grad) = torch.autograd.grad(
            [conv_copy_out], [inp_copy, weight_copy, bias_copy], [grad_copy]
        )

        # check numerics
        torch.testing.assert_close(bias_copy_grad.cpu(), bias_grad.cpu())

        torch.testing.assert_close(weight_copy_grad.cpu(), weight_grad.cpu())
        torch.testing.assert_close(inp_copy_grad.cpu(), inp_grad.cpu())
//  [python]
def test_adjust_sharpness(self):
        x_shape = [4, 4, 3]
        x_data = [75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 0,
                  0, 65, 108, 101, 120, 97, 110, 100, 101, 114, 32, 86, 114, 121, 110, 105,
                  111, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117]
        x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
        x_pil = Image.fromarray(x_np, mode='RGB')

        # test 0
        y_pil = F.adjust_sharpness(x_pil, 1)
        y_np = np.array(y_pil)
        self.assertEqual(y_np.tolist(), x_np.tolist())

        # test 1
        y_pil = F.adjust_sharpness(x_pil, 0.5)
        y_np = np.array(y_pil)
        y_ans = [75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 30,
                 30, 74, 103, 96, 114, 97, 110, 100, 101, 114, 32, 81, 103, 108, 102, 101,
                 107, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117]
        y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
        self.assertEqual(y_np.tolist(), y_ans.tolist())

        # test 2
        y_pil = F.adjust_sharpness(x_pil, 2)
        y_np = np.array(y_pil)
        y_ans = [75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 0,
                 0, 46, 118, 111, 132, 97, 110, 100, 101, 114, 32, 95, 135, 146, 126, 112,
                 119, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117]
        y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
        self.assertEqual(y_np.tolist(), y_ans.tolist())

        # test 3
        x_shape = [2, 2, 3]
        x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
        x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
        x_pil = Image.fromarray(x_np, mode='RGB')
        x_th = torch.tensor(x_np.transpose(2, 0, 1))
        y_pil = F.adjust_sharpness(x_pil, 2)
        y_np = np.array(y_pil).transpose(2, 0, 1)
        y_th = F.adjust_sharpness(x_th, 2)
        self.assertEqual(y_np.tolist(), y_th.numpy().tolist())
//  [python]
def test_adjust_sharpness(self):
        x_shape = [4, 4, 3]
        x_data = [75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 0,
                  0, 65, 108, 101, 120, 97, 110, 100, 101, 114, 32, 86, 114, 121, 110, 105,
                  111, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117]
        x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
        x_pil = Image.fromarray(x_np, mode='RGB')

        # test 0
        y_pil = F.adjust_sharpness(x_pil, 1)
        y_np = np.array(y_pil)
        torch.testing.assert_close(y_np, x_np)

        # test 1
        y_pil = F.adjust_sharpness(x_pil, 0.5)
        y_np = np.array(y_pil)
        y_ans = [75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 30,
                 30, 74, 103, 96, 114, 97, 110, 100, 101, 114, 32, 81, 103, 108, 102, 101,
                 107, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117]
        y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
        torch.testing.assert_close(y_np, y_ans)

        # test 2
        y_pil = F.adjust_sharpness(x_pil, 2)
        y_np = np.array(y_pil)
        y_ans = [75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 0,
                 0, 46, 118, 111, 132, 97, 110, 100, 101, 114, 32, 95, 135, 146, 126, 112,
                 119, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117]
        y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
        torch.testing.assert_close(y_np, y_ans)

        # test 3
        x_shape = [2, 2, 3]
        x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
        x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
        x_pil = Image.fromarray(x_np, mode='RGB')
        x_th = torch.tensor(x_np.transpose(2, 0, 1))
        y_pil = F.adjust_sharpness(x_pil, 2)
        y_np = np.array(y_pil).transpose(2, 0, 1)
        y_th = F.adjust_sharpness(x_th, 2)
        torch.testing.assert_close(y_np, y_th.numpy())
//  [python]
def test_sample_all():
    for c, val, _, group_circuit in make_discrete_var_test_circuits():
        c.print()
        assert c.shape == val.shape[1:]
        sampled_circ = Sampler(RunDiscreteVarAllSpec([group_circuit])).estimate_and_sample(c)
        transform = sampled_circ.evaluate()
        t.assert_close(transform, val, atol=1e-5, rtol=1e-5)
//  [python]
def test_sample_all():
    for c, val, _, group_circuit in make_discrete_var_test_circuits():
        c.print()
        assert c.shape == val.shape[1:]
        sampled_circ = Sampler(RunDiscreteVarAllSpec([group_circuit])).estimate_and_sample(c)
        transform = sampled_circ.evaluate()
        torch.testing.assert_close(transform, val, atol=1e-5, rtol=1e-5)
//  [python]
def test_seed_per_circuit_sample(sample_shape):
    u, _, _ = make_discrete_var((3, 42), 17, "u")
    v, _, _ = make_discrete_var((3, 42), 17, "v")

    def seeder(c):
        if c == u.probs_and_group:
            return 0
        if c == v.probs_and_group:
            return 1
        return hash(c) ^ 2833

    c1 = Add(v, u)
    a1 = Sampler(RandomSampleSpec(sample_shape, seeder=seeder)).estimate_and_sample(c1).evaluate()
    c2 = Add(u, v)
    a2 = Sampler(RandomSampleSpec(sample_shape, seeder=seeder)).estimate_and_sample(c2).evaluate()
    t.assert_close(a1, a2)
//  [python]
def test_seed_per_circuit_sample(sample_shape):
    u, _, _ = make_discrete_var((3, 42), 17, "u")
    v, _, _ = make_discrete_var((3, 42), 17, "v")

    def seeder(c):
        if c == u.probs_and_group:
            return 0
        if c == v.probs_and_group:
            return 1
        return hash(c) ^ 2833

    c1 = Add(v, u)
    a1 = Sampler(RandomSampleSpec(sample_shape, seeder=seeder)).estimate_and_sample(c1).evaluate()
    c2 = Add(u, v)
    a2 = Sampler(RandomSampleSpec(sample_shape, seeder=seeder)).estimate_and_sample(c2).evaluate()
    torch.testing.assert_close(a1, a2)
//  [python]
def test_seed_sample(sample_shape):
    for v, _, *_ in make_discrete_var_test_circuits():
        # Check no error
        orig_shape = v.shape
        torch.manual_seed(1234)
        a1 = Sampler(RandomSampleSpec(sample_shape)).estimate_and_sample(v).evaluate()
        assert a1.shape == sample_shape + orig_shape
        torch.manual_seed(1234)
        a2 = Sampler(RandomSampleSpec(sample_shape)).estimate_and_sample(v).evaluate()
        torch.manual_seed(1234)
        a3 = Sampler(RandomSampleSpec(sample_shape, seeder=None)).estimate_and_sample(v).evaluate()
        t.assert_close(a1, a2)
        t.assert_close(a1, a3)
        t.assert_close(a2, a3)

        a1 = Sampler(RandomSampleSpec(sample_shape, seeder=default_hash_seeder(87))).estimate_and_sample(v).evaluate()
        assert a1.shape == sample_shape + orig_shape
        a2 = Sampler(RandomSampleSpec(sample_shape, seeder=default_hash_seeder(87))).estimate_and_sample(v).evaluate()
        t.assert_close(a1, a2)
//  [python]
def test_seed_sample(sample_shape):
    for v, _, *_ in make_discrete_var_test_circuits():
        # Check no error
        orig_shape = v.shape
        torch.manual_seed(1234)
        a1 = Sampler(RandomSampleSpec(sample_shape)).estimate_and_sample(v).evaluate()
        assert a1.shape == sample_shape + orig_shape
        torch.manual_seed(1234)
        a2 = Sampler(RandomSampleSpec(sample_shape)).estimate_and_sample(v).evaluate()
        torch.manual_seed(1234)
        a3 = Sampler(RandomSampleSpec(sample_shape, seeder=None)).estimate_and_sample(v).evaluate()
        torch.testing.assert_close(a1, a2)
        torch.testing.assert_close(a1, a3)
        torch.testing.assert_close(a2, a3)

        a1 = Sampler(RandomSampleSpec(sample_shape, seeder=default_hash_seeder(87))).estimate_and_sample(v).evaluate()
        assert a1.shape == sample_shape + orig_shape
        a2 = Sampler(RandomSampleSpec(sample_shape, seeder=default_hash_seeder(87))).estimate_and_sample(v).evaluate()
        torch.testing.assert_close(a1, a2)
//  [python]
def test_transformer_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,
        )
        # Transformer greedy doesn't return attention scores
        # `return_attention = False` by default
        self.assertIsNone(attention_scores)

        # outputs
        self.assertEqual(output.shape, (batch_size, max_output_length))  # batch x time
        self.assertEqual(output, self.expected_transformer_ids)

        # scores
        self.assertEqual(scores.shape, (batch_size, max_output_length))  # batch x time
        self.assertEqual(scores, self.expected_transformer_scores)
//  [python]
def test_transformer_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,
        )
        # Transformer greedy doesn't return attention scores
        # `return_attention = False` by default
        self.assertIsNone(attention_scores)

        # outputs
        self.assertEqual(output.shape, (batch_size, max_output_length))  # batch x time
        torch.testing.assert_close(
            output, self.expected_transformer_ids, check_dtype=False
        )

        # scores
        self.assertEqual(scores.shape, (batch_size, max_output_length))  # batch x time
        torch.testing.assert_close(
            scores, self.expected_transformer_scores, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_transformer_beam1(self):
        batch_size = 2
        beam_size = 1
        alpha = 0.0
        n_best = 1
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        beam_output, beam_scores, attention_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,
        )
        # Transformer beam doesn't return attention scores
        self.assertIsNone(attention_scores)

        # batch_size * n_best x hyp_len
        self.assertEqual(beam_output.shape, (batch_size * n_best, max_output_length))
        self.assertEqual(beam_output.dtype, self.expected_transformer_ids.dtype)
        torch.testing.assert_allclose(
            beam_output, self.expected_transformer_ids, rtol=1e-4, atol=1e-4
        )
        torch.testing.assert_allclose(
            beam_scores, torch.tensor([[-1.5772], [-2.8292]]), rtol=1e-4, atol=1e-4
        )

        # now compare to greedy, they should be the same for beam=1
        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(beam_output.dtype, greedy_output.dtype)
        torch.testing.assert_allclose(beam_output, greedy_output, rtol=1e-4, atol=1e-4)
        torch.testing.assert_allclose(
            greedy_scores, self.expected_transformer_scores, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_transformer_beam1(self):
        batch_size = 2
        beam_size = 1
        alpha = 0.0
        n_best = 1
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        beam_output, beam_scores, attention_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,
        )
        # Transformer beam doesn't return attention scores
        self.assertIsNone(attention_scores)

        # batch_size * n_best x hyp_len
        self.assertEqual(beam_output.shape, (batch_size * n_best, max_output_length))
        torch.testing.assert_close(
            beam_output, self.expected_transformer_ids, check_dtype=False
        )
        torch.testing.assert_close(
            beam_scores, torch.tensor([[-1.5772], [-2.8292]]), rtol=1e-4, atol=1e-4
        )

        # now compare to greedy, they should be the same for beam=1
        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,
        )
        torch.testing.assert_close(beam_output, greedy_output, check_dtype=False)
        torch.testing.assert_close(
            greedy_scores, self.expected_transformer_scores, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_ngram_blocker(self):
        batch_size = 2
        max_output_length = 7
        no_repeat_ngram_size = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        output, scores, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            encoder_input=None,
            return_prob="hyp",
            no_repeat_ngram_size=no_repeat_ngram_size,
            autocast=self.autocat,
        )

        expected_output = torch.tensor([[0, 0, 0, 0, 0, 0, 0], [0, 0, 7, 0, 1, 0, 0]])
        torch.allclose(output, expected_output, check_dtype=False)

        expected_scores = torch.tensor([
            [-0.5425, -0.4908, -0.5439, -0.7328, -0.6922, -0.6422, -0.6464],
            [-0.8726, -0.9898, -0.9668, -1.3988, -0.6783, -1.0269, -0.6804],
        ])
        torch.allclose(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def test_ngram_blocker(self):
        batch_size = 2
        max_output_length = 7
        no_repeat_ngram_size = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        output, scores, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            encoder_input=None,
            return_prob="hyp",
            no_repeat_ngram_size=no_repeat_ngram_size,
            autocast=self.autocat,
        )

        expected_output = torch.tensor([[0, 0, 0, 0, 0, 0, 0], [0, 0, 7, 0, 1, 0, 0]])
        torch.testing.assert_close(output, expected_output, check_dtype=False)

        expected_scores = torch.tensor([
            [-0.5425, -0.4908, -0.5439, -0.7328, -0.6922, -0.6422, -0.6464],
            [-0.8726, -0.9898, -0.9668, -1.3988, -0.6783, -1.0269, -0.6804],
        ])
        torch.testing.assert_close(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def test_transformer_beam7(self):
        batch_size = 2
        beam_size = 7
        n_best = 5
        alpha = 1.0
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        output, scores, attention_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,
        )
        # Transformer beam doesn't return attention scores
        self.assertIsNone(attention_scores)

        # batch_size * n_best x hyp_len (= time steps)
        self.assertEqual(output.shape, (batch_size * n_best, max_output_length))

        expected_output = torch.tensor([
            [0, 0, 0], [0, 0, 7], [0, 7, 0], [7, 0, 0], [7, 0, 7],
            [0, 0, 7], [7, 0, 0], [0, 0, 0], [0, 7, 0], [7, 0, 7],
        ])  # yapf: disable
        torch.testing.assert_allclose(output, expected_output)

        expected_scores = torch.tensor([
            [-1.1829], [-1.6948], [-1.7128], [-2.0805], [-2.9899],
            [-2.1219], [-2.1881], [-2.1931], [-2.3707], [-2.4195],
        ])  # yapf: disable
        torch.testing.assert_allclose(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def test_transformer_beam7(self):
        batch_size = 2
        beam_size = 7
        n_best = 5
        alpha = 1.0
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        output, scores, attention_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,
        )
        # Transformer beam doesn't return attention scores
        self.assertIsNone(attention_scores)

        # batch_size * n_best x hyp_len (= time steps)
        self.assertEqual(output.shape, (batch_size * n_best, max_output_length))

        expected_output = torch.tensor([
            [0, 0, 0], [0, 0, 7], [0, 7, 0], [7, 0, 0], [7, 0, 7],
            [0, 0, 7], [7, 0, 0], [0, 0, 0], [0, 7, 0], [7, 0, 7],
        ])  # yapf: disable
        torch.testing.assert_close(output, expected_output)

        expected_scores = torch.tensor([
            [-1.1829], [-1.6948], [-1.7128], [-2.0805], [-2.9899],
            [-2.1219], [-2.1881], [-2.1931], [-2.3707], [-2.4195],
        ])  # yapf: disable
        torch.testing.assert_close(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def test_repetition_penalty_in_beam_search(self):
        batch_size = 2
        beam_size = 7
        n_best = 5
        alpha = 1.0
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        src_tokens = torch.tensor([[5, 5, 4], [5, 6, 6]]).long()
        output_with_penalty, scores_with_penalty, _ = 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,
            encoder_input=src_tokens,
            repetition_penalty=1.5,
            return_prob="hyp",
            autocast=self.autocat,
        )

        expected_output_with_penalty = torch.tensor([
            [0, 0, 0], [0, 7, 0], [0, 0, 7], [7, 0, 0], [7, 0, 7],
            [7, 0, 0], [0, 0, 7], [7, 0, 7], [0, 7, 0], [0, 0, 0],
        ])  # yapf: disable
        t.assert_close(
            output_with_penalty, expected_output_with_penalty, check_dtype=False
        )

        expected_scores_with_penalty = torch.tensor([
            [-1.5709], [-1.8617], [-1.8788], [-2.2284], [-3.5925],
            [-2.4791], [-2.4931], [-2.8261], [-2.8357], [-2.9624],
        ])  # yapf: disable
        t.assert_close(
            scores_with_penalty, expected_scores_with_penalty, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_repetition_penalty_in_beam_search(self):
        batch_size = 2
        beam_size = 7
        n_best = 5
        alpha = 1.0
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        src_tokens = torch.tensor([[5, 5, 4], [5, 6, 6]]).long()
        output_with_penalty, scores_with_penalty, _ = 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,
            encoder_input=src_tokens,
            repetition_penalty=1.5,
            return_prob="hyp",
            autocast=self.autocat,
        )

        expected_output_with_penalty = torch.tensor([
            [0, 0, 0], [0, 7, 0], [0, 0, 7], [7, 0, 0], [7, 0, 7],
            [7, 0, 0], [0, 0, 7], [7, 0, 7], [0, 7, 0], [0, 0, 0],
        ])  # yapf: disable
        torch.testing.assert_close(
            output_with_penalty, expected_output_with_penalty, check_dtype=False
        )

        expected_scores_with_penalty = torch.tensor([
            [-1.5709], [-1.8617], [-1.8788], [-2.2284], [-3.5925],
            [-2.4791], [-2.4931], [-2.8261], [-2.8357], [-2.9624],
        ])  # yapf: disable
        torch.testing.assert_close(
            scores_with_penalty, expected_scores_with_penalty, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_correct_cumulant_estimate(seed=283838):
    torch.manual_seed(seed)
    x = torch.randn(5, 3, 2, dtype=torch.float64)
    y = torch.nn.functional.softmax(x, dim=-1).sum(dim=1)
    z = x ** 3
    w = torch.sigmoid(x[:, :, 0])

    vals = [x, y, z, w]

    var_x = DiscreteVar.new_uniform(Array(x))
    vars_unweighted = [DiscreteVar(Array(v), var_x.probs_and_group) for v in vals]

    probs = torch.rand(5, dtype=torch.float64)
    probs = probs / probs.sum()

    probs_and_group = Array(probs)

    vars_weighted = [DiscreteVar(Array(v), probs_and_group) for v in vals]

    def estimate_and_comp(circuit, actual):
        # todo: add is_full to rust
        transformed = (
            Sampler(
                RunDiscreteVarAllSpec.create_full_from_circuits(circuit),
            )
            .estimate(circuit)
            .evaluate()
        )
        t.assert_close(
            transformed,
            actual,
            atol=1e-6,
            rtol=1e-6,
        )

    for all_vars, weights in [
        [vars_unweighted, torch.full((5,), 1 / 5, dtype=torch.float64)],
        [vars_weighted, probs],
    ]:
        var_x, var_y, var_z, var_w = all_vars
        all_means = [torch.einsum("i ..., i -> ...", v, weights).unsqueeze(0) for v in vals]
        c_x, c_y, c_z, c_w = [v - m for v, m in zip(vals, all_means)]
        x_mean = all_means[0].squeeze(0)
        estimate_and_comp(
            Cumulant(
                var_x,
            ),
            x_mean,
        )

        uncorrected_cov = torch.einsum("b i j, b k, b -> i j k", c_x, c_y, weights)
        estimate_and_comp(Cumulant(var_x, var_y), uncorrected_cov)

        uncorrected_third_cum = torch.einsum("b i j, b k, b l m, b -> i j k l m", c_x, c_y, c_z, weights)
        estimate_and_comp(Cumulant(var_x, var_y, var_z), uncorrected_third_cum)

        centered_4th_mom = torch.einsum("b i j, b k, b l m, b o, b -> i j k l m o", c_x, c_y, c_z, c_w, weights)
        first_pair = torch.einsum(
            "i j k, l m o -> i j k l m o",
            torch.einsum("b i j, b k, b -> i j k", c_x, c_y, weights),
            torch.einsum("b l m, b o, b -> l m o", c_z, c_w, weights),
        )
        snd_pair = torch.einsum(
            "i j l m, k o -> i j k l m o",
            torch.einsum("b i j, b l m, b -> i j l m", c_x, c_z, weights),
            torch.einsum("b k, b o, b-> k o", c_y, c_w, weights),
        )
        thrd_pair = torch.einsum(
            "i j o, k l m -> i j k l m o",
            torch.einsum("b i j, b o, b -> i j o", c_x, c_w, weights),
            torch.einsum("b k, b l m, b -> k l m", c_y, c_z, weights),
        )

        pair_sum = first_pair + snd_pair + thrd_pair

        uncorrected_4th_cum = centered_4th_mom - pair_sum

        estimate_and_comp(Cumulant(var_x, var_y, var_z, var_w), uncorrected_4th_cum)
//  [python]
def test_correct_cumulant_estimate(seed=283838):
    torch.manual_seed(seed)
    x = torch.randn(5, 3, 2, dtype=torch.float64)
    y = torch.nn.functional.softmax(x, dim=-1).sum(dim=1)
    z = x ** 3
    w = torch.sigmoid(x[:, :, 0])

    vals = [x, y, z, w]

    var_x = DiscreteVar.new_uniform(Array(x))
    vars_unweighted = [DiscreteVar(Array(v), var_x.probs_and_group) for v in vals]

    probs = torch.rand(5, dtype=torch.float64)
    probs = probs / probs.sum()

    probs_and_group = Array(probs)

    vars_weighted = [DiscreteVar(Array(v), probs_and_group) for v in vals]

    def estimate_and_comp(circuit, actual):
        # todo: add is_full to rust
        transformed = (
            Sampler(
                RunDiscreteVarAllSpec.create_full_from_circuits(circuit),
            )
            .estimate(circuit)
            .evaluate()
        )
        torch.testing.assert_close(
            transformed,
            actual,
            atol=1e-6,
            rtol=1e-6,
        )

    for all_vars, weights in [
        [vars_unweighted, torch.full((5,), 1 / 5, dtype=torch.float64)],
        [vars_weighted, probs],
    ]:
        var_x, var_y, var_z, var_w = all_vars
        all_means = [torch.einsum("i ..., i -> ...", v, weights).unsqueeze(0) for v in vals]
        c_x, c_y, c_z, c_w = [v - m for v, m in zip(vals, all_means)]
        x_mean = all_means[0].squeeze(0)
        estimate_and_comp(
            Cumulant(
                var_x,
            ),
            x_mean,
        )

        uncorrected_cov = torch.einsum("b i j, b k, b -> i j k", c_x, c_y, weights)
        estimate_and_comp(Cumulant(var_x, var_y), uncorrected_cov)

        uncorrected_third_cum = torch.einsum("b i j, b k, b l m, b -> i j k l m", c_x, c_y, c_z, weights)
        estimate_and_comp(Cumulant(var_x, var_y, var_z), uncorrected_third_cum)

        centered_4th_mom = torch.einsum("b i j, b k, b l m, b o, b -> i j k l m o", c_x, c_y, c_z, c_w, weights)
        first_pair = torch.einsum(
            "i j k, l m o -> i j k l m o",
            torch.einsum("b i j, b k, b -> i j k", c_x, c_y, weights),
            torch.einsum("b l m, b o, b -> l m o", c_z, c_w, weights),
        )
        snd_pair = torch.einsum(
            "i j l m, k o -> i j k l m o",
            torch.einsum("b i j, b l m, b -> i j l m", c_x, c_z, weights),
            torch.einsum("b k, b o, b-> k o", c_y, c_w, weights),
        )
        thrd_pair = torch.einsum(
            "i j o, k l m -> i j k l m o",
            torch.einsum("b i j, b o, b -> i j o", c_x, c_w, weights),
            torch.einsum("b k, b l m, b -> k l m", c_y, c_z, weights),
        )

        pair_sum = first_pair + snd_pair + thrd_pair

        uncorrected_4th_cum = centered_4th_mom - pair_sum

        estimate_and_comp(Cumulant(var_x, var_y, var_z, var_w), uncorrected_4th_cum)
//  [python]
def test_ngram_blocker_in_beam_search(self):
        batch_size = 2
        beam_size = 3
        n_best = 3
        alpha = 1.0
        max_output_length = 7
        no_repeat_ngram_size = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        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,
            encoder_input=None,
            return_prob="hyp",
            no_repeat_ngram_size=no_repeat_ngram_size,
            autocast=self.autocat,
        )

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

        expected_scores = torch.tensor([
            [-2.1454], [-2.4287], [-2.4680], [-3.2931], [-3.3489], [-3.4080]
        ])  # yapf: disable
        torch.testing.assert_allclose(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
def test_ngram_blocker_in_beam_search(self):
        batch_size = 2
        beam_size = 3
        n_best = 3
        alpha = 1.0
        max_output_length = 7
        no_repeat_ngram_size = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        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,
            encoder_input=None,
            return_prob="hyp",
            no_repeat_ngram_size=no_repeat_ngram_size,
            autocast=self.autocat,
        )

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

        expected_scores = torch.tensor([
            [-2.1454], [-2.4287], [-2.4680], [-3.2931], [-3.3489], [-3.4080]
        ])  # yapf: disable
        torch.testing.assert_close(scores, expected_scores, rtol=1e-4, atol=1e-4)
//  [python]
input = torch.randn(1)
mod = Foo()
compiled_mod = compiled_module(mod, fw_compiler, bw_compiler)

for a, b in zip(run(mod, input), run(compiled_mod, input)):
    t.assert_close(a, b)
//  [python]
input = torch.randn(1)
mod = Foo()
compiled_mod = compiled_module(mod, fw_compiler, bw_compiler)

for a, b in zip(run(mod, input), run(compiled_mod, input)):
    torch.testing.assert_close(a, b)
//  [python]
def test_transformer_greedy_with_prompt(self):
        batch_size = 2
        max_output_length = 7
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        decoder_prompt = torch.tensor([[2, 7, 7, 4], [0, 7, 4, 1]])
        trg_prompt_mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]])
        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",
            return_attention=True,
            autocast=self.autocat,
            decoder_prompt=decoder_prompt,
            trg_prompt_mask=trg_prompt_mask,
        )

        expected_output = torch.tensor([[7, 7, 4, 0, 0, 0, 0], [7, 4, 0, 0, 0, 0, 0]])
        # forced decoding
        self.assertEqual(output.shape, (batch_size, max_output_length))  # batch x time
        self.assertEqual(output, expected_output)

        # zero log_prob on the forced positions
        expected_score = torch.tensor([
            [0.0000, 0.0000, 0.0000, -0.4631, -0.3289, -0.3042, -0.3404],
            [0.0000, 0.0000, -0.7532, -0.6751, -0.5730, -0.4957, -0.5718],
        ])
        self.assertEqual(scores.shape, (batch_size, max_output_length))  # batch x time
        self.assertTrue(torch.allclose(scores, expected_score, rtol=1e-4, atol=1e-4))

        expected_att = torch.tensor([
            [[0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000, 0.0000],
             [0.0000, 0.0000, 0.0000, 0.0000], [0.3926, 0.2844, 0.3191, 0.0039],
             [0.4019, 0.2798, 0.3156, 0.0027], [0.4072, 0.2809, 0.3093, 0.0026],
             [0.4004, 0.2799, 0.3169, 0.0027]],
            [[0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000, 0.0000],
             [0.3194, 0.0042, 0.4271, 0.2492], [0.3523, 0.0036, 0.3957, 0.2484],
             [0.3335, 0.0034, 0.4143, 0.2488], [0.3135, 0.0031, 0.4346, 0.2488],
             [0.3322, 0.0034, 0.4158, 0.2486]],
        ])
        self.assertEqual(
            attention_scores.shape,  # batch x trg_len x src_len
            (batch_size, max_output_length, encoder_output.size(1))
        )
        self.assertTrue(torch.allclose(attention_scores, expected_att, rtol=1e-4, atol=1e-4))
//  [python]
def test_transformer_greedy_with_prompt(self):
        batch_size = 2
        max_output_length = 7
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        decoder_prompt = torch.tensor([[2, 7, 7, 4], [0, 7, 4, 1]])
        trg_prompt_mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]])
        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",
            return_attention=True,
            autocast=self.autocat,
            decoder_prompt=decoder_prompt,
            trg_prompt_mask=trg_prompt_mask,
        )

        expected_output = torch.tensor([[7, 7, 4, 0, 0, 0, 0], [7, 4, 0, 0, 0, 0, 0]])
        # forced decoding
        self.assertEqual(output.shape, (batch_size, max_output_length))  # batch x time
        torch.testing.assert_close(output, expected_output, check_dtype=False)

        # zero log_prob on the forced positions
        expected_score = torch.tensor([
            [0.0000, 0.0000, 0.0000, -0.4631, -0.3289, -0.3042, -0.3404],
            [0.0000, 0.0000, -0.7532, -0.6751, -0.5730, -0.4957, -0.5718],
        ])
        self.assertEqual(scores.shape, (batch_size, max_output_length))  # batch x time
        torch.testing.assert_close(scores, expected_score, rtol=1e-4, atol=1e-4)

        expected_att = torch.tensor([
            [[0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000, 0.0000],
             [0.0000, 0.0000, 0.0000, 0.0000], [0.3926, 0.2844, 0.3191, 0.0039],
             [0.4019, 0.2798, 0.3156, 0.0027], [0.4072, 0.2809, 0.3093, 0.0026],
             [0.4004, 0.2799, 0.3169, 0.0027]],
            [[0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000, 0.0000],
             [0.3194, 0.0042, 0.4271, 0.2492], [0.3523, 0.0036, 0.3957, 0.2484],
             [0.3335, 0.0034, 0.4143, 0.2488], [0.3135, 0.0031, 0.4346, 0.2488],
             [0.3322, 0.0034, 0.4158, 0.2486]],
        ])
        self.assertEqual(
            attention_scores.shape,  # batch x trg_len x src_len
            (batch_size, max_output_length, encoder_output.size(1))
        )
        torch.testing.assert_close(attention_scores, expected_att, rtol=1e-4, atol=1e-4)
//  [python]
def test_repetition_penalty_and_generate_unk(self):
        batch_size = 3
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        # no repetition penalty
        output, _, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            generate_unk=False,
            autocast=self.autocat,
        )

        expected_output = torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
        self.assertTrue(torch.allclose(output, expected_output, check_dtype=False))
        self.assertEqual(torch.count_nonzero(output).item(), 9)  # no unk token

        # trg repetition penalty
        output_trg_penalty, _, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            encoder_input=None,
            repetition_penalty=1.5,
            generate_unk=False,
            autocast=self.autocat,
        )

        expected_output_trg_penalty = torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 3]])
        self.assertTrue(torch.allclose(
            output_trg_penalty, expected_output_trg_penalty, check_dtype=False
        ))

        # src + trg repetition penalty
        # src_len = 4 (see self._build())
        src_tokens = torch.tensor([[4, 3, 1, 1], [5, 4, 3, 1], [5, 5, 6, 3]]).long()
        src_mask = (src_tokens != 1).unsqueeze(1)
        output_src_penalty, _, attention = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            encoder_input=src_tokens,
            repetition_penalty=1.5,
            generate_unk=False,
            return_attention=True,
            autocast=self.autocat,
        )

        expected_output_src_penalty = torch.tensor([[1, 7, 3], [1, 7, 1], [1, 1, 1]])
        self.assertTrue(torch.allclose(
            output_src_penalty, expected_output_src_penalty, check_dtype=False
        ))

        # Transformer Greedy can return attention probs
        # (batch_size, trg_len, src_len) = (3, 3, 4)
        expected_attention = torch.tensor([
            [[0.5292, 0.4708, 0.0000, 0.0000], [0.5269, 0.4731, 0.0000, 0.0000],
             [0.5264, 0.4736, 0.0000, 0.0000]],
            [[0.3075, 0.6322, 0.0602, 0.0000], [0.2890, 0.6350, 0.0760, 0.0000],
             [0.3343, 0.6314, 0.0343, 0.0000]],
            [[0.2648, 0.1326, 0.5174, 0.0852], [0.2642, 0.1167, 0.5365, 0.0825],
             [0.2646, 0.1125, 0.5421, 0.0809]],
        ])
        self.assertTrue(torch.allclose(attention, expected_attention, rtol=1e-4, atol=1e-4))
//  [python]
def test_repetition_penalty_and_generate_unk(self):
        batch_size = 3
        max_output_length = 3
        src_mask, model, encoder_output, encoder_hidden = self._build(
            batch_size=batch_size
        )

        # no repetition penalty
        output, _, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            generate_unk=False,
            autocast=self.autocat,
        )

        expected_output = torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
        torch.testing.assert_close(output, expected_output, check_dtype=False)
        self.assertEqual(torch.count_nonzero(output).item(), 9)  # no unk token

        # trg repetition penalty
        output_trg_penalty, _, _ = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            encoder_input=None,
            repetition_penalty=1.5,
            generate_unk=False,
            autocast=self.autocat,
        )

        expected_output_trg_penalty = torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 3]])
        torch.testing.assert_close(
            output_trg_penalty, expected_output_trg_penalty, check_dtype=False
        )

        # src + trg repetition penalty
        # src_len = 4 (see self._build())
        src_tokens = torch.tensor([[4, 3, 1, 1], [5, 4, 3, 1], [5, 5, 6, 3]]).long()
        src_mask = (src_tokens != 1).unsqueeze(1)
        output_src_penalty, _, attention = greedy(
            src_mask=src_mask,
            max_output_length=max_output_length,
            model=model,
            encoder_output=encoder_output,
            encoder_hidden=encoder_hidden,
            encoder_input=src_tokens,
            repetition_penalty=1.5,
            generate_unk=False,
            return_attention=True,
            autocast=self.autocat,
        )

        expected_output_src_penalty = torch.tensor([[1, 7, 3], [1, 7, 1], [1, 1, 1]])
        torch.testing.assert_close(
            output_src_penalty, expected_output_src_penalty, check_dtype=False
        )

        # Transformer Greedy can return attention probs
        # (batch_size, trg_len, src_len) = (3, 3, 4)
        expected_attention = torch.tensor([
            [[0.5292, 0.4708, 0.0000, 0.0000], [0.5269, 0.4731, 0.0000, 0.0000],
             [0.5264, 0.4736, 0.0000, 0.0000]],
            [[0.3075, 0.6322, 0.0602, 0.0000], [0.2890, 0.6350, 0.0760, 0.0000],
             [0.3343, 0.6314, 0.0343, 0.0000]],
            [[0.2648, 0.1326, 0.5174, 0.0852], [0.2642, 0.1167, 0.5365, 0.0825],
             [0.2646, 0.1125, 0.5421, 0.0809]],
        ])
        torch.testing.assert_close(attention, expected_attention, rtol=1e-4, atol=1e-4)
//  [python]
def test_interpolate(self):

    arg = (torch.randn(3, 3, 200, 200),)
    model = Interpolate()

    ans = model(*arg)

    with torch.no_grad():
      exported = torch.export.export(model, arg)
      shlo = exported_program_to_stablehlo(exported)
      ans2 = shlo(*arg).cpu().to(torch.float32)
      torch.testing.assert_allclose(ans, ans2, rtol=1e-5, atol=1e-4)
//  [python]
def test_interpolate(self):

    arg = (torch.randn(3, 3, 200, 200),)
    model = Interpolate()

    ans = model(*arg)

    with torch.no_grad():
      exported = torch.export.export(model, arg)
      shlo = exported_program_to_stablehlo(exported)
      ans2 = shlo(*arg).cpu().to(torch.float32)
      torch.testing.assert_close(ans, ans2, rtol=1e-5, atol=1e-4)
//  [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))
        self.assertTrue(torch.allclose(greedy_output, self.expected_recurrent_ids, check_dtype=False))
        self.assertTrue(torch.allclose(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,
        )

        self.assertTrue(torch.allclose(greedy_output, beam_output, check_dtype=False))
        self.assertTrue(torch.allclose(beam_scores, torch.tensor([[-4.7406], [-5.0405]]), rtol=1e-4, atol=1e-4))