jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_int_scaled_mm(self, device, dtype):
        if device == "cuda" and not torch.cuda.is_available():
            self.skipTest(f"{device} not available")

        from torchao.kernel import intmm

        dtype = torch.bfloat16
        m, k, n = (128, 64, 16)
        x = torch.randn(m, k, dtype=dtype, device=device)
        scales = x.sum(-1, keepdim=True)
        w = torch.randn(n, k, dtype=dtype, device=device).t()
        x_int = x.to(dtype=torch.int8)
        w_int = w.to(dtype=torch.int8)
        out32_1 = intmm.safe_int_mm(x_int, w_int) * scales
        assert out32_1.dtype == torch.bfloat16
        out32_2 = intmm.int_scaled_matmul(x_int, w_int, scales)
        assert out32_2.dtype == out32_1.dtype
        torch.testing.assert_allclose(out32_1, out32_2)
//  [python]
def test_int_scaled_mm(self, device, dtype):
        if device == "cuda" and not torch.cuda.is_available():
            self.skipTest(f"{device} not available")

        from torchao.kernel import intmm

        dtype = torch.bfloat16
        m, k, n = (128, 64, 16)
        x = torch.randn(m, k, dtype=dtype, device=device)
        scales = x.sum(-1, keepdim=True)
        w = torch.randn(n, k, dtype=dtype, device=device).t()
        x_int = x.to(dtype=torch.int8)
        w_int = w.to(dtype=torch.int8)
        out32_1 = intmm.safe_int_mm(x_int, w_int) * scales
        assert out32_1.dtype == torch.bfloat16
        out32_2 = intmm.int_scaled_matmul(x_int, w_int, scales)
        assert out32_2.dtype == out32_1.dtype
        torch.testing.assert_close(out32_1, out32_2)
//  [python]
def test_cholesky_inverse_no_batch():
    """Checks that our Cholesky inverse matches `torch.cholesky_inverse()`."""
    torch.autograd.set_detect_anomaly(True)
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(fannypack.utils.tril_count_from_matrix_dim(matrix_dim))
    )
    torch.testing.assert_allclose(
        fannypack.utils.cholesky_inverse(L), torch.cholesky_inverse(L)
    )
//  [python]
def test_cholesky_inverse_no_batch():
    """Checks that our Cholesky inverse matches `torch.cholesky_inverse()`."""
    torch.autograd.set_detect_anomaly(True)
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(fannypack.utils.tril_count_from_matrix_dim(matrix_dim))
    )
    torch.testing.assert_close(
        fannypack.utils.cholesky_inverse(L), torch.cholesky_inverse(L)
    )
//  [python]
def test_cholesky_inverse():
    """Checks that our Cholesky inverse matches `torch.cholesky_inverse()`."""
    torch.autograd.set_detect_anomaly(True)
    batch_dims = (5,)
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),),
        )
    )

    for i, our_inverse in enumerate(fannypack.utils.cholesky_inverse(L)):
        torch.testing.assert_allclose(our_inverse, torch.cholesky_inverse(L[i]))
//  [python]
def test_cholesky_inverse():
    """Checks that our Cholesky inverse matches `torch.cholesky_inverse()`."""
    torch.autograd.set_detect_anomaly(True)
    batch_dims = (5,)
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),),
        )
    )

    for i, our_inverse in enumerate(fannypack.utils.cholesky_inverse(L)):
        torch.testing.assert_close(our_inverse, torch.cholesky_inverse(L[i]))
//  [python]
def test_tril_inverse():
    """Check that our `tril_inverse` function correctly inverts some full-rank
    matrices.
    """
    for matrix_dim in range(2, 5):
        tril_matrix = fannypack.utils.tril_from_vector(
            torch.randn((5, fannypack.utils.tril_count_from_matrix_dim(matrix_dim)))
        )
        inverse = fannypack.utils.tril_inverse(tril_matrix)

        torch.testing.assert_allclose(tril_matrix @ inverse, torch.eye(matrix_dim))
//  [python]
def test_tril_inverse():
    """Check that our `tril_inverse` function correctly inverts some full-rank
    matrices.
    """
    for matrix_dim in range(2, 5):
        tril_matrix = fannypack.utils.tril_from_vector(
            torch.randn((5, fannypack.utils.tril_count_from_matrix_dim(matrix_dim)))
        )
        inverse = fannypack.utils.tril_inverse(tril_matrix)

        torch.testing.assert_close(tril_matrix @ inverse, torch.eye(matrix_dim))
//  [python]
def test_cholupdate():
    """Checks rank-1 Cholesky update forward pass."""
    batch_dims = (1,)  # (5, 3, 7)
    matrix_dim = 2
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),)
        )
    )
    x = torch.randn(batch_dims + (matrix_dim,))

    updated_L = fannypack.utils.cholupdate(L, x)

    torch.testing.assert_allclose(
        L @ L.transpose(-1, -2) + x[..., :, None] @ x[..., None, :],
        updated_L @ updated_L.transpose(-1, -2),
    )
//  [python]
def test_cholupdate():
    """Checks rank-1 Cholesky update forward pass."""
    batch_dims = (1,)  # (5, 3, 7)
    matrix_dim = 2
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),)
        )
    )
    x = torch.randn(batch_dims + (matrix_dim,))

    updated_L = fannypack.utils.cholupdate(L, x)

    torch.testing.assert_close(
        L @ L.transpose(-1, -2) + x[..., :, None] @ x[..., None, :],
        updated_L @ updated_L.transpose(-1, -2),
    )
//  [python]
def test_empty_reference(self):
        best_logp, best_alig = ctc_alignment(self._logp, [], ctc_sym=0)
        torch.testing.assert_allclose(best_logp, math.log(0.3 * 0.4 * 0.5 * 0.1))
        self.assertEqual([0, 0, 0, 0], best_alig)
//  [python]
def test_empty_reference(self):
        best_logp, best_alig = ctc_alignment(self._logp, [], ctc_sym=0)
        torch.testing.assert_close(best_logp, math.log(0.3 * 0.4 * 0.5 * 0.1))
        self.assertEqual([0, 0, 0, 0], best_alig)
//  [python]
def test_empty_reference2(self):
        best_logp, best_alig = ctc_alignment(self._logp, [], ctc_sym=2)
        torch.testing.assert_allclose(best_logp, math.log(0.2 * 0.1 * 0.4 * 0.2))
        self.assertEqual([2, 2, 2, 2], best_alig)
//  [python]
def test_empty_reference2(self):
        best_logp, best_alig = ctc_alignment(self._logp, [], ctc_sym=2)
        torch.testing.assert_close(best_logp, math.log(0.2 * 0.1 * 0.4 * 0.2))
        self.assertEqual([2, 2, 2, 2], best_alig)
//  [python]
def test_single_label(self):
        best_logp, best_alig = ctc_alignment(self._logp, [1], ctc_sym=0)
        torch.testing.assert_allclose(best_logp, math.log(0.3 * 0.4 * 0.5 * 0.7))
        self.assertEqual([0, 0, 0, 1], best_alig)
//  [python]
def test_single_label(self):
        best_logp, best_alig = ctc_alignment(self._logp, [1], ctc_sym=0)
        torch.testing.assert_close(best_logp, math.log(0.3 * 0.4 * 0.5 * 0.7))
        self.assertEqual([0, 0, 0, 1], best_alig)
//  [python]
def test_single_label2(self):
        best_logp, best_alig = ctc_alignment(self._logp, [2], ctc_sym=1)
        torch.testing.assert_allclose(best_logp, math.log(0.5 * 0.5 * 0.4 * 0.7))
        self.assertEqual([1, 1, 2, 1], best_alig)
//  [python]
def test_single_label2(self):
        best_logp, best_alig = ctc_alignment(self._logp, [2], ctc_sym=1)
        torch.testing.assert_close(best_logp, math.log(0.5 * 0.5 * 0.4 * 0.7))
        self.assertEqual([1, 1, 2, 1], best_alig)
//  [python]
def test_repeated_label(self):
        best_logp, best_alig = ctc_alignment(self._logp, [2, 2], ctc_sym=0)
        torch.testing.assert_allclose(best_logp, math.log(0.2 * 0.4 * 0.5 * 0.2))
        self.assertEqual([2, 0, 0, 2], best_alig)
//  [python]
def test_repeated_label(self):
        best_logp, best_alig = ctc_alignment(self._logp, [2, 2], ctc_sym=0)
        torch.testing.assert_close(best_logp, math.log(0.2 * 0.4 * 0.5 * 0.2))
        self.assertEqual([2, 0, 0, 2], best_alig)
//  [python]
def test_example(self, discrete, dim_state, dim_action, kind, dim_reward):
        if discrete:
            num_states, num_actions = dim_state, dim_action
            dim_state, dim_action = (), ()
        else:
            num_states, num_actions = -1, -1
            dim_state, dim_action = (dim_state,), (dim_action,)
        if kind == "nan":
            o = Observation.nan_example(
                dim_state=dim_state,
                dim_action=dim_action,
                num_states=num_states,
                num_actions=num_actions,
                dim_reward=dim_reward,
            )
        elif kind == "zero":
            o = Observation.zero_example(
                dim_state=dim_state,
                dim_action=dim_action,
                num_states=num_states,
                num_actions=num_actions,
                dim_reward=dim_reward,
            )
        elif kind == "random":
            o = Observation.random_example(
                dim_state=dim_state,
                dim_action=dim_action,
                num_states=num_states,
                num_actions=num_actions,
                dim_reward=dim_reward,
            )
        else:
            with pytest.raises(ValueError):
                Observation.get_example(
                    dim_state=dim_state,
                    dim_action=dim_action,
                    num_states=num_states,
                    num_actions=num_actions,
                    dim_reward=dim_reward,
                    kind=kind,
                )
            return

        if discrete:
            torch.testing.assert_allclose(o.state.shape, torch.Size([]))
            torch.testing.assert_allclose(o.action.shape, torch.Size([]))
            torch.testing.assert_allclose(o.next_state.shape, torch.Size([]))
            torch.testing.assert_allclose(o.log_prob_action, torch.tensor(1.0))

        else:
            torch.testing.assert_allclose(o.state.shape, torch.Size(dim_state))
            torch.testing.assert_allclose(o.action.shape, torch.Size(dim_action))
            torch.testing.assert_allclose(o.next_state.shape, torch.Size(dim_state))
            torch.testing.assert_allclose(o.log_prob_action, torch.tensor(1.0))
//  [python]
def test_example(self, discrete, dim_state, dim_action, kind, dim_reward):
        if discrete:
            num_states, num_actions = dim_state, dim_action
            dim_state, dim_action = (), ()
        else:
            num_states, num_actions = -1, -1
            dim_state, dim_action = (dim_state,), (dim_action,)
        if kind == "nan":
            o = Observation.nan_example(
                dim_state=dim_state,
                dim_action=dim_action,
                num_states=num_states,
                num_actions=num_actions,
                dim_reward=dim_reward,
            )
        elif kind == "zero":
            o = Observation.zero_example(
                dim_state=dim_state,
                dim_action=dim_action,
                num_states=num_states,
                num_actions=num_actions,
                dim_reward=dim_reward,
            )
        elif kind == "random":
            o = Observation.random_example(
                dim_state=dim_state,
                dim_action=dim_action,
                num_states=num_states,
                num_actions=num_actions,
                dim_reward=dim_reward,
            )
        else:
            with pytest.raises(ValueError):
                Observation.get_example(
                    dim_state=dim_state,
                    dim_action=dim_action,
                    num_states=num_states,
                    num_actions=num_actions,
                    dim_reward=dim_reward,
                    kind=kind,
                )
            return

        if discrete:
            torch.testing.assert_close(o.state.shape, torch.Size([]))
            torch.testing.assert_close(o.action.shape, torch.Size([]))
            torch.testing.assert_close(o.next_state.shape, torch.Size([]))
            torch.testing.assert_close(o.log_prob_action, torch.tensor(1.0))

        else:
            torch.testing.assert_close(o.state.shape, torch.Size(dim_state))
            torch.testing.assert_close(o.action.shape, torch.Size(dim_action))
            torch.testing.assert_close(o.next_state.shape, torch.Size(dim_state))
            torch.testing.assert_close(o.log_prob_action, torch.tensor(1.0))
//  [python]
def test_regular(self):
        best_logp, best_alig = ctc_alignment(self._logp, [1, 2, 2], ctc_sym=0)
        torch.testing.assert_allclose(best_logp, math.log(0.5 * 0.1 * 0.5 * 0.2))
        self.assertEqual([1, 2, 0, 2], best_alig)
//  [python]
def test_regular(self):
        best_logp, best_alig = ctc_alignment(self._logp, [1, 2, 2], ctc_sym=0)
        torch.testing.assert_close(best_logp, math.log(0.5 * 0.1 * 0.5 * 0.2))
        self.assertEqual([1, 2, 0, 2], best_alig)
//  [python]
def test_regular2(self):
        best_logp, best_alig = ctc_alignment(self._logp, [1, 0], ctc_sym=2)
        torch.testing.assert_allclose(best_logp, math.log(0.5 * 0.5 * 0.5 * 0.2))
        self.assertEqual([1, 1, 0, 2], best_alig)
//  [python]
def test_regular2(self):
        best_logp, best_alig = ctc_alignment(self._logp, [1, 0], ctc_sym=2)
        torch.testing.assert_close(best_logp, math.log(0.5 * 0.5 * 0.5 * 0.2))
        self.assertEqual([1, 1, 0, 2], best_alig)
//  [python]
def test_regular3(self):
        best_logp, best_alig = ctc_alignment(self._logp, [0, 2, 0, 2], ctc_sym=1)
        torch.testing.assert_allclose(best_logp, math.log(0.3 * 0.1 * 0.5 * 0.2))
        self.assertEqual([0, 2, 0, 2], best_alig)
//  [python]
def test_regular3(self):
        best_logp, best_alig = ctc_alignment(self._logp, [0, 2, 0, 2], ctc_sym=1)
        torch.testing.assert_close(best_logp, math.log(0.3 * 0.1 * 0.5 * 0.2))
        self.assertEqual([0, 2, 0, 2], best_alig)
//  [python]
def test_alias_analysis_inputs(self):
        class AliasModule(nn.Module):
            def __init__(self):
                super(AliasModule, self).__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                x.add_(y)
                w = z + self.a
                z = w + x
                return z

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am
        am = getModule(False)
        am_s = getModule(True)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        ref = am(x, x, x)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        test = am_s(x, x, x)

        torch.testing.assert_allclose(ref, test)
//  [python]
def test_alias_analysis_inputs(self):
        class AliasModule(nn.Module):
            def __init__(self):
                super(AliasModule, self).__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                x.add_(y)
                w = z + self.a
                z = w + x
                return z

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am
        am = getModule(False)
        am_s = getModule(True)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        ref = am(x, x, x)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        test = am_s(x, x, x)

        torch.testing.assert_close(ref, test)
//  [python]
def test_alias_analysis_module(self):
        class AliasModule(nn.Module):
            def __init__(self):
                super(AliasModule, self).__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                z = z + self.a
                self.b.add_(y)
                w = z + self.a
                z = w + x
                return z
        x = torch.randn(128, 128)

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am

        am = getModule(False)
        am_s = getModule(True)
        ref = am(x, x, x)
        test = am_s(x, x, x)
        torch.testing.assert_allclose(ref, test)

        # Now do the aliasing
        am.a = am.b
        ref = am(x, x, x)

        am_s.a = am_s.b
        test = am_s(x, x, x)

        torch.testing.assert_allclose(ref, test)
//  [python]
def test_alias_analysis_module(self):
        class AliasModule(nn.Module):
            def __init__(self):
                super(AliasModule, self).__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                z = z + self.a
                self.b.add_(y)
                w = z + self.a
                z = w + x
                return z
        x = torch.randn(128, 128)

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am

        am = getModule(False)
        am_s = getModule(True)
        ref = am(x, x, x)
        test = am_s(x, x, x)
        torch.testing.assert_close(ref, test)

        # Now do the aliasing
        am.a = am.b
        ref = am(x, x, x)

        am_s.a = am_s.b
        test = am_s(x, x, x)

        torch.testing.assert_close(ref, test)
//  [python]
def test_alias_analysis_input_and_module(self):
        class AliasModule(nn.Module):
            def __init__(self):
                super(AliasModule, self).__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                x.add_(y)
                w = z + self.b
                z = w + x
                return z

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am
        am = getModule(False)
        am_s = getModule(True)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        am.b = x
        ref = am(x, x, x)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        am_s.b = x
        test = am_s(x, x, x)

        torch.testing.assert_allclose(ref, test)
//  [python]
def test_alias_analysis_input_and_module(self):
        class AliasModule(nn.Module):
            def __init__(self):
                super(AliasModule, self).__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                x.add_(y)
                w = z + self.b
                z = w + x
                return z

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am
        am = getModule(False)
        am_s = getModule(True)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        am.b = x
        ref = am(x, x, x)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        am_s.b = x
        test = am_s(x, x, x)

        torch.testing.assert_close(ref, test)
//  [python]
def test_cle_conv(self):
        class TestModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv0 = nn.Conv2d(3, 8, 2)
                self.conv1 = nn.Conv2d(8, 16, 2)
                self.conv2 = nn.Conv2d(16, 64, 2)

            def forward(self, x):
                conv0 = self.conv0(x)
                conv1 = self.conv1(conv0)
                conv2 = self.conv2(conv1)
                return conv2

        torch.manual_seed(10)

        dummy_input = torch.randn(1, 3, 224, 224)
        model = TestModel()
        model.eval()

        origin_output = model(dummy_input)
        cle_model = cross_layer_equalize(model, dummy_input, torch.device('cpu'), hba_flag=False)
        cle_output = cle_model(dummy_input)
        torch.testing.assert_allclose(origin_output, cle_output)
//  [python]
def test_cle_conv(self):
        class TestModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv0 = nn.Conv2d(3, 8, 2)
                self.conv1 = nn.Conv2d(8, 16, 2)
                self.conv2 = nn.Conv2d(16, 64, 2)

            def forward(self, x):
                conv0 = self.conv0(x)
                conv1 = self.conv1(conv0)
                conv2 = self.conv2(conv1)
                return conv2

        torch.manual_seed(10)

        dummy_input = torch.randn(1, 3, 224, 224)
        model = TestModel()
        model.eval()

        origin_output = model(dummy_input)
        cle_model = cross_layer_equalize(model, dummy_input, torch.device('cpu'), hba_flag=False)
        cle_output = cle_model(dummy_input)
        torch.testing.assert_close(origin_output, cle_output)
//  [python]
def test_invert():
    qs = torch.tensor([5, 3, 3, 2, 0])
    ps = invert_indices(qs, max(qs)+1)

    im, ma = ps.kth(qs, 0)

    torch.testing.assert_allclose(qs, qs[im])
//  [python]
def test_invert():
    qs = torch.tensor([5, 3, 3, 2, 0])
    ps = invert_indices(qs, max(qs)+1)

    im, ma = ps.kth(qs, 0)

    torch.testing.assert_close(qs, qs[im])
//  [python]
def test_cle_group_conv(self):
        class TestModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv0 = nn.Conv2d(3, 8, 2)
                self.conv1 = nn.Conv2d(8, 8, 2, groups=4)
                self.conv1 = nn.Conv2d(8, 8, 2, groups=8)
                self.conv2 = nn.Conv2d(8, 16, 2)

            def forward(self, x):
                conv0 = self.conv0(x)
                conv1 = self.conv1(conv0)
                conv2 = self.conv2(conv1)
                return conv2

        torch.manual_seed(10)

        dummy_input = torch.randn(1, 3, 224, 224)
        model = TestModel()
        model.eval()

        origin_output = model(dummy_input)
        cle_model = cross_layer_equalize(model, dummy_input, torch.device('cpu'), hba_flag=False)
        cle_output = cle_model(dummy_input)
        torch.testing.assert_allclose(origin_output, cle_output)
//  [python]
def test_cle_group_conv(self):
        class TestModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv0 = nn.Conv2d(3, 8, 2)
                self.conv1 = nn.Conv2d(8, 8, 2, groups=4)
                self.conv1 = nn.Conv2d(8, 8, 2, groups=8)
                self.conv2 = nn.Conv2d(8, 16, 2)

            def forward(self, x):
                conv0 = self.conv0(x)
                conv1 = self.conv1(conv0)
                conv2 = self.conv2(conv1)
                return conv2

        torch.manual_seed(10)

        dummy_input = torch.randn(1, 3, 224, 224)
        model = TestModel()
        model.eval()

        origin_output = model(dummy_input)
        cle_model = cross_layer_equalize(model, dummy_input, torch.device('cpu'), hba_flag=False)
        cle_output = cle_model(dummy_input)
        torch.testing.assert_close(origin_output, cle_output)
//  [python]
def assert_equal(x, y):
    if isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor):
        torch.testing.assert_allclose(x, y.to(x.device))
    elif isinstance(x, OrderedDict) and isinstance(y, OrderedDict):
        for x_value, y_value in zip(x.values(), y.values()):
            assert_equal(x_value, y_value)
    elif isinstance(x, dict) and isinstance(y, dict):
        assert x.keys() == y.keys()
        for key in x.keys():
            assert_equal(x[key], y[key])
    elif isinstance(x, str) and isinstance(y, str):
        assert x == y
    elif isinstance(x, Iterable) and isinstance(y, Iterable):
        assert len(x) == len(y)
        for x_item, y_item in zip(x, y):
            assert_equal(x_item, y_item)
    else:
        assert x == y
//  [python]
def assert_equal(x, y):
    if isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor):
        torch.testing.assert_close(x, y.to(x.device))
    elif isinstance(x, OrderedDict) and isinstance(y, OrderedDict):
        for x_value, y_value in zip(x.values(), y.values()):
            assert_equal(x_value, y_value)
    elif isinstance(x, dict) and isinstance(y, dict):
        assert x.keys() == y.keys()
        for key in x.keys():
            assert_equal(x[key], y[key])
    elif isinstance(x, str) and isinstance(y, str):
        assert x == y
    elif isinstance(x, Iterable) and isinstance(y, Iterable):
        assert len(x) == len(y)
        for x_item, y_item in zip(x, y):
            assert_equal(x_item, y_item)
    else:
        assert x == y
//  [python]
def test_fast_init_context_manager(self):
        # 1. Create a dummy class. Should have buffers as well? To make sure we test __init__
        class MyClass(PreTrainedModel):
            config_class = PretrainedConfig

            def __init__(self, config=None):
                super().__init__(config if config is not None else PretrainedConfig())
                self.linear = nn.Linear(10, 10, bias=True)
                self.embedding = nn.Embedding(10, 10)
                self.std = 1

            def _init_weights(self, module):
                if isinstance(module, nn.Linear):
                    module.weight.data = nn.init.kaiming_uniform_(module.weight.data, np.sqrt(5))
                    if module.bias is not None:
                        module.bias.data.normal_(mean=0.0, std=self.std)

        # 2. Make sure a linear layer's reset params is properly skipped:
        with ContextManagers([no_init_weights(True)]):
            no_init_instance = MyClass()

        set_seed(0)
        expected_bias = torch.tensor(
            ([0.2975, 0.2131, -0.1379, -0.0796, -0.3012, -0.0057, -0.2381, -0.2439, -0.0174, 0.0475])
        )
        init_instance = MyClass()
        torch.testing.assert_allclose(init_instance.linear.bias, expected_bias, rtol=1e-3, atol=1e-4)

        set_seed(0)
        torch.testing.assert_allclose(
            init_instance.linear.weight, nn.init.kaiming_uniform_(no_init_instance.linear.weight, np.sqrt(5))
        )

        # 3. Make sure weights that are not present use init_weight_ and get expected values
        with tempfile.TemporaryDirectory() as tmpdirname:
            state_dict = init_instance.state_dict()
            del state_dict["linear.weight"]

            init_instance.config.save_pretrained(tmpdirname)
            torch.save(state_dict, os.path.join(tmpdirname, "pytorch_model.bin"))
            set_seed(0)
            model_fast_init = MyClass.from_pretrained(tmpdirname)

            set_seed(0)
            model_slow_init = MyClass.from_pretrained(tmpdirname, _fast_init=False)

            for key in model_fast_init.state_dict().keys():
                max_diff = torch.max(torch.abs(model_slow_init.state_dict()[key] - model_fast_init.state_dict()[key]))
                self.assertLessEqual(max_diff.item(), 1e-3, msg=f"{key} not identical")
//  [python]
def test_fast_init_context_manager(self):
        # 1. Create a dummy class. Should have buffers as well? To make sure we test __init__
        class MyClass(PreTrainedModel):
            config_class = PretrainedConfig

            def __init__(self, config=None):
                super().__init__(config if config is not None else PretrainedConfig())
                self.linear = nn.Linear(10, 10, bias=True)
                self.embedding = nn.Embedding(10, 10)
                self.std = 1

            def _init_weights(self, module):
                if isinstance(module, nn.Linear):
                    module.weight.data = nn.init.kaiming_uniform_(module.weight.data, np.sqrt(5))
                    if module.bias is not None:
                        module.bias.data.normal_(mean=0.0, std=self.std)

        # 2. Make sure a linear layer's reset params is properly skipped:
        with ContextManagers([no_init_weights(True)]):
            no_init_instance = MyClass()

        set_seed(0)
        expected_bias = torch.tensor(
            ([0.2975, 0.2131, -0.1379, -0.0796, -0.3012, -0.0057, -0.2381, -0.2439, -0.0174, 0.0475])
        )
        init_instance = MyClass()
        torch.testing.assert_close(init_instance.linear.bias, expected_bias, rtol=1e-3, atol=1e-4)

        set_seed(0)
        torch.testing.assert_close(
            init_instance.linear.weight, nn.init.kaiming_uniform_(no_init_instance.linear.weight, np.sqrt(5))
        )

        # 3. Make sure weights that are not present use init_weight_ and get expected values
        with tempfile.TemporaryDirectory() as tmpdirname:
            state_dict = init_instance.state_dict()
            del state_dict["linear.weight"]

            init_instance.config.save_pretrained(tmpdirname)
            torch.save(state_dict, os.path.join(tmpdirname, "pytorch_model.bin"))
            set_seed(0)
            model_fast_init = MyClass.from_pretrained(tmpdirname)

            set_seed(0)
            model_slow_init = MyClass.from_pretrained(tmpdirname, _fast_init=False)

            for key in model_fast_init.state_dict().keys():
                max_diff = torch.max(torch.abs(model_slow_init.state_dict()[key] - model_fast_init.state_dict()[key]))
                self.assertLessEqual(max_diff.item(), 1e-3, msg=f"{key} not identical")
//  [python]
def test_conv_dec_streaming(channel, kernel_size, stride):
    input_audio = torch.randn((1, 100))
    ilens = torch.LongTensor([100])

    encoder = ConvEncoder(kernel_size=kernel_size, stride=stride, channel=channel)
    decoder = ConvDecoder(kernel_size=kernel_size, stride=stride, channel=channel)
    frames, flens = encoder(input_audio, ilens)
    wav, ilens = decoder(frames, ilens)

    splited = encoder.streaming_frame(input_audio)

    sframes = [encoder.forward_streaming(s) for s in splited]
    swavs = [decoder.forward_streaming(s) for s in sframes]
    merged = decoder.streaming_merge(swavs, ilens)

    sframes = torch.cat(sframes, dim=1)

    torch.testing.assert_allclose(sframes, frames)
    torch.testing.assert_allclose(wav, merged)
//  [python]
def test_conv_dec_streaming(channel, kernel_size, stride):
    input_audio = torch.randn((1, 100))
    ilens = torch.LongTensor([100])

    encoder = ConvEncoder(kernel_size=kernel_size, stride=stride, channel=channel)
    decoder = ConvDecoder(kernel_size=kernel_size, stride=stride, channel=channel)
    frames, flens = encoder(input_audio, ilens)
    wav, ilens = decoder(frames, ilens)

    splited = encoder.streaming_frame(input_audio)

    sframes = [encoder.forward_streaming(s) for s in splited]
    swavs = [decoder.forward_streaming(s) for s in sframes]
    merged = decoder.streaming_merge(swavs, ilens)

    sframes = torch.cat(sframes, dim=1)

    torch.testing.assert_close(sframes, frames)
    torch.testing.assert_close(wav, merged)
//  [python]
def test_inference_long_sentence(self):
        tokenizer = FNetTokenizerFast.from_pretrained("google/fnet-base")

        inputs = tokenizer(
            "the man worked as a [MASK].",
            "this is his [MASK].",
            return_tensors="pt",
            padding="max_length",
            max_length=512,
        )

        torch.testing.assert_allclose(inputs["input_ids"], torch.tensor([[4, 13, 283, 2479, 106, 8, 6, 845, 5, 168, 65, 367, 6, 845, 5, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3]]))  # fmt: skip

        inputs = {k: v.to(torch_device) for k, v in inputs.items()}

        model = FNetForMaskedLM.from_pretrained("google/fnet-base")
        model.to(torch_device)
        logits = model(**inputs).logits
        predictions_mask_1 = tokenizer.decode(logits[0, 6].topk(5).indices)
        predictions_mask_2 = tokenizer.decode(logits[0, 12].topk(5).indices)

        self.assertEqual(predictions_mask_1.split(" "), ["man", "child", "teacher", "woman", "model"])
        self.assertEqual(predictions_mask_2.split(" "), ["work", "wife", "job", "story", "name"])
//  [python]
def test_inference_long_sentence(self):
        tokenizer = FNetTokenizerFast.from_pretrained("google/fnet-base")

        inputs = tokenizer(
            "the man worked as a [MASK].",
            "this is his [MASK].",
            return_tensors="pt",
            padding="max_length",
            max_length=512,
        )

        torch.testing.assert_close(inputs["input_ids"], torch.tensor([[4, 13, 283, 2479, 106, 8, 6, 845, 5, 168, 65, 367, 6, 845, 5, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3]]))  # fmt: skip

        inputs = {k: v.to(torch_device) for k, v in inputs.items()}

        model = FNetForMaskedLM.from_pretrained("google/fnet-base")
        model.to(torch_device)
        logits = model(**inputs).logits
        predictions_mask_1 = tokenizer.decode(logits[0, 6].topk(5).indices)
        predictions_mask_2 = tokenizer.decode(logits[0, 12].topk(5).indices)

        self.assertEqual(predictions_mask_1.split(" "), ["man", "child", "teacher", "woman", "model"])
        self.assertEqual(predictions_mask_2.split(" "), ["work", "wife", "job", "story", "name"])
//  [python]
def test_shape_and_type(self, rewards, gamma):
        np_returns = discount_cumsum(np.atleast_2d(np.array(rewards)).T, gamma)
        assert np_returns.shape == (len(rewards), 1)
        assert type(np_returns) is np.ndarray

        t_returns = discount_cumsum(
            torch.tensor(rewards, dtype=torch.get_default_dtype()).unsqueeze(-1), gamma
        )
        assert t_returns.shape == torch.Size((len(rewards), 1))
        assert type(t_returns) is torch.Tensor

        torch.testing.assert_allclose(t_returns, np_returns)
//  [python]
def test_shape_and_type(self, rewards, gamma):
        np_returns = discount_cumsum(np.atleast_2d(np.array(rewards)).T, gamma)
        assert np_returns.shape == (len(rewards), 1)
        assert type(np_returns) is np.ndarray

        t_returns = discount_cumsum(
            torch.tensor(rewards, dtype=torch.get_default_dtype()).unsqueeze(-1), gamma
        )
        assert t_returns.shape == torch.Size((len(rewards), 1))
        assert type(t_returns) is torch.Tensor

        torch.testing.assert_close(t_returns, np_returns)
//  [python]
def test_tensor(self):
        x = torch.randn(3, 5, 7, 9, requires_grad=True)
        layer = MaskImageFromSize(mask_value=-99)
        y = layer(x)
        torch.testing.assert_allclose(y, x)
        (dx,) = torch.autograd.grad([torch.sum(y)], [x])
        torch.testing.assert_allclose(dx, torch.ones(3, 5, 7, 9))
//  [python]
def test_tensor(self):
        x = torch.randn(3, 5, 7, 9, requires_grad=True)
        layer = MaskImageFromSize(mask_value=-99)
        y = layer(x)
        torch.testing.assert_close(y, x)
        (dx,) = torch.autograd.grad([torch.sum(y)], [x])
        torch.testing.assert_close(dx, torch.ones(3, 5, 7, 9))
//  [python]
def test_correctness(self, gamma, batch):
        rewards = np.array([[1, 2], [0.5, 0.3], [2, -1.2], [-0.2, 0.5]])

        cum_rewards = np.array(
            [
                [
                    1 + 0.5 * gamma + 2 * gamma ** 2 - 0.2 * gamma ** 3,
                    2 + 0.3 * gamma - 1.2 * gamma ** 2 + 0.5 * gamma ** 3,
                ],
                [
                    0.5 + 2 * gamma - 0.2 * gamma ** 2,
                    0.3 - 1.2 * gamma + 0.5 * gamma ** 2,
                ],
                [2 - 0.2 * gamma, -1.2 + 0.5 * gamma],
                [-0.2, 0.5],
            ]
        )
        if batch:
            rewards = np.tile(np.array(rewards), (5, 1, 1))
            cum_rewards = np.tile(np.array(cum_rewards), (5, 1, 1))
        assert scipy.allclose(cum_rewards, discount_cumsum(np.array(rewards), gamma))

        torch.testing.assert_allclose(
            torch.tensor(cum_rewards), discount_cumsum(torch.tensor(rewards), gamma)
        )
        torch.testing.assert_allclose(
            torch.tensor(cum_rewards[..., 0, :], dtype=torch.get_default_dtype()),
            discount_sum(torch.tensor(rewards, dtype=torch.get_default_dtype()), gamma),
        )

        for i in range(rewards.shape[-1]):
            torch.testing.assert_allclose(
                torch.tensor(cum_rewards)[..., [i]],
                discount_cumsum(torch.tensor(rewards)[..., [i]], gamma),
            )
            torch.testing.assert_allclose(
                torch.tensor(cum_rewards[..., 0, [i]], dtype=torch.get_default_dtype()),
                discount_sum(
                    torch.tensor(rewards[..., [i]], dtype=torch.get_default_dtype()),
                    gamma,
                ),
            )
//  [python]
def test_correctness(self, gamma, batch):
        rewards = np.array([[1, 2], [0.5, 0.3], [2, -1.2], [-0.2, 0.5]])

        cum_rewards = np.array(
            [
                [
                    1 + 0.5 * gamma + 2 * gamma ** 2 - 0.2 * gamma ** 3,
                    2 + 0.3 * gamma - 1.2 * gamma ** 2 + 0.5 * gamma ** 3,
                ],
                [
                    0.5 + 2 * gamma - 0.2 * gamma ** 2,
                    0.3 - 1.2 * gamma + 0.5 * gamma ** 2,
                ],
                [2 - 0.2 * gamma, -1.2 + 0.5 * gamma],
                [-0.2, 0.5],
            ]
        )
        if batch:
            rewards = np.tile(np.array(rewards), (5, 1, 1))
            cum_rewards = np.tile(np.array(cum_rewards), (5, 1, 1))
        assert scipy.allclose(cum_rewards, discount_cumsum(np.array(rewards), gamma))

        torch.testing.assert_close(
            torch.tensor(cum_rewards), discount_cumsum(torch.tensor(rewards), gamma)
        )
        torch.testing.assert_close(
            torch.tensor(cum_rewards[..., 0, :], dtype=torch.get_default_dtype()),
            discount_sum(torch.tensor(rewards, dtype=torch.get_default_dtype()), gamma),
        )

        for i in range(rewards.shape[-1]):
            torch.testing.assert_close(
                torch.tensor(cum_rewards)[..., [i]],
                discount_cumsum(torch.tensor(rewards)[..., [i]], gamma),
            )
            torch.testing.assert_close(
                torch.tensor(cum_rewards[..., 0, [i]], dtype=torch.get_default_dtype()),
                discount_sum(
                    torch.tensor(rewards[..., [i]], dtype=torch.get_default_dtype()),
                    gamma,
                ),
            )
//  [python]
def test_acc(xlist, blist, rtol=1e-3, atol=1e-5, tolerate_small_mismatch=True):
    for a, b in zip(xlist, blist):
        try:
            torch.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
        except AssertionError as error:
            if tolerate_small_mismatch:
                print("small mismatch detected", error)
            else:
                return False
    return True
//  [python]
def test_acc(xlist, blist, rtol=1e-3, atol=1e-5, tolerate_small_mismatch=True):
    for a, b in zip(xlist, blist):
        try:
            torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
        except AssertionError as error:
            if tolerate_small_mismatch:
                print("small mismatch detected", error)
            else:
                return False
    return True
//  [python]
def test_run(self):
        command = [
            'python',
            'tools/analysis_tools/confusion_matrix.py',
            str(self.config_file),
            str(self.result_file),
            '--out',
            str(self.dir / 'result.pkl'),
        ]
        Popen(command, cwd=MMPRE_ROOT, stdout=PIPE).wait()
        result = mmengine.load(self.dir / 'result.pkl')
        torch.testing.assert_allclose(
            result, torch.tensor([
                [1, 0, 0],
                [1, 0, 0],
                [0, 0, 0],
            ]))
//  [python]
def test_run(self):
        command = [
            'python',
            'tools/analysis_tools/confusion_matrix.py',
            str(self.config_file),
            str(self.result_file),
            '--out',
            str(self.dir / 'result.pkl'),
        ]
        Popen(command, cwd=MMPRE_ROOT, stdout=PIPE).wait()
        result = mmengine.load(self.dir / 'result.pkl')
        torch.testing.assert_close(
            result, torch.tensor([
                [1, 0, 0],
                [1, 0, 0],
                [0, 0, 0],
            ]))
//  [python]
def test_correctness(self, gamma, value_function, entropy_reg):
        trajectory = [
            Observation(0, 0, reward=np.array([1]), done=False, entropy=0.2).to_torch(),
            Observation(
                0, 0, reward=np.array([0.5]), done=False, entropy=0.3
            ).to_torch(),
            Observation(0, 0, reward=np.array([2]), done=False, entropy=0.5).to_torch(),
            Observation(
                0, 0, reward=np.array([-0.2]), done=False, entropy=-0.2
            ).to_torch(),
        ]

        r0 = 1 + entropy_reg * 0.2
        r1 = 0.5 + entropy_reg * 0.3
        r2 = 2 + entropy_reg * 0.5
        r3 = -0.2 - entropy_reg * 0.2

        v = 0.01 if value_function is not None else 0

        reward = mc_return(
            stack_list_of_tuples(trajectory, 0),
            gamma,
            value_function=value_function,
            entropy_regularization=entropy_reg,
            reduction="min",
        )

        torch.testing.assert_allclose(
            reward,
            torch.tensor(
                [r0 + r1 * gamma + r2 * gamma ** 2 + r3 * gamma ** 3 + v * gamma ** 4]
            ),
        )
        torch.testing.assert_allclose(
            mc_return(
                observation=Observation(state=0, reward=np.array([0])).to_torch(),
                gamma=gamma,
                value_function=value_function,
                entropy_regularization=entropy_reg,
            ),
            torch.tensor([0]),
        )
//  [python]
def test_correctness(self, gamma, value_function, entropy_reg):
        trajectory = [
            Observation(0, 0, reward=np.array([1]), done=False, entropy=0.2).to_torch(),
            Observation(
                0, 0, reward=np.array([0.5]), done=False, entropy=0.3
            ).to_torch(),
            Observation(0, 0, reward=np.array([2]), done=False, entropy=0.5).to_torch(),
            Observation(
                0, 0, reward=np.array([-0.2]), done=False, entropy=-0.2
            ).to_torch(),
        ]

        r0 = 1 + entropy_reg * 0.2
        r1 = 0.5 + entropy_reg * 0.3
        r2 = 2 + entropy_reg * 0.5
        r3 = -0.2 - entropy_reg * 0.2

        v = 0.01 if value_function is not None else 0

        reward = mc_return(
            stack_list_of_tuples(trajectory, 0),
            gamma,
            value_function=value_function,
            entropy_regularization=entropy_reg,
            reduction="min",
        )

        torch.testing.assert_close(
            reward,
            torch.tensor(
                [r0 + r1 * gamma + r2 * gamma ** 2 + r3 * gamma ** 3 + v * gamma ** 4]
            ),
        )
        torch.testing.assert_close(
            mc_return(
                observation=Observation(state=0, reward=np.array([0])).to_torch(),
                gamma=gamma,
                value_function=value_function,
                entropy_regularization=entropy_reg,
            ),
            torch.tensor([0]),
        )
//  [python]
def test_padded_tensor(self):
        x = torch.randn(3, 5, 7, 9, requires_grad=True)
        xs = torch.tensor([[1, 1], [3, 9], [7, 5]])
        layer = MaskImageFromSize(mask_value=-99)
        y = layer(PaddedTensor(x, xs))
        # Expected output
        expected_y = x.clone()
        expected_y[0, :, 1:, :] = -99
        expected_y[0, :, :, 1:] = -99
        expected_y[1, :, 3:, :] = -99
        expected_y[2, :, :, 5:] = -99
        # Check expected output
        torch.testing.assert_allclose(y.data, expected_y)

        # Test backward pass (mask with 0, so that we can sum without masking)
        layer = MaskImageFromSize(mask_value=0)
        y = layer(PaddedTensor(x, xs))
        (dx,) = torch.autograd.grad([torch.sum(y.data)], [x])
        # Expected output
        expected_dx = torch.ones(3, 5, 7, 9)
        expected_dx[0, :, 1:, :] = 0
        expected_dx[0, :, :, 1:] = 0
        expected_dx[1, :, 3:, :] = 0
        expected_dx[2, :, :, 5:] = 0
        # Check expected output
        torch.testing.assert_allclose(dx, expected_dx)
//  [python]
def test_padded_tensor(self):
        x = torch.randn(3, 5, 7, 9, requires_grad=True)
        xs = torch.tensor([[1, 1], [3, 9], [7, 5]])
        layer = MaskImageFromSize(mask_value=-99)
        y = layer(PaddedTensor(x, xs))
        # Expected output
        expected_y = x.clone()
        expected_y[0, :, 1:, :] = -99
        expected_y[0, :, :, 1:] = -99
        expected_y[1, :, 3:, :] = -99
        expected_y[2, :, :, 5:] = -99
        # Check expected output
        torch.testing.assert_close(y.data, expected_y)

        # Test backward pass (mask with 0, so that we can sum without masking)
        layer = MaskImageFromSize(mask_value=0)
        y = layer(PaddedTensor(x, xs))
        (dx,) = torch.autograd.grad([torch.sum(y.data)], [x])
        # Expected output
        expected_dx = torch.ones(3, 5, 7, 9)
        expected_dx[0, :, 1:, :] = 0
        expected_dx[0, :, :, 1:] = 0
        expected_dx[1, :, 3:, :] = 0
        expected_dx[2, :, :, 5:] = 0
        # Check expected output
        torch.testing.assert_close(dx, expected_dx)
//  [python]
def test_init(memory):
    torch.testing.assert_allclose(memory.next_done, torch.tensor([0, 0, 0, 0]))
    assert memory.next_value is None
    assert memory.device == torch.device("cpu")
    assert memory.global_step == 0
    assert memory.obs_preprocessor is not None
    assert memory.experiences == []
    assert memory.episode_lengths == []
    assert memory.episode_returns == []
    assert memory.vars_to_log == {}