jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_layer_norm(train_config: TrainConfig, elementwise_affine: bool, include_bias: bool):
    train_config.model.layer_norm_with_affine = elementwise_affine
    train_config.model.include_bias = include_bias
    ln = LayerNorm.build(train_config.model)

    needs_weight = elementwise_affine
    needs_bias = elementwise_affine and include_bias
    with torch.no_grad():
        if needs_weight:
            weight = torch.randn(train_config.model.d_model)
            ln.weight.copy_(weight)
        else:
            weight = None

        if needs_bias:
            bias = torch.randn(train_config.model.d_model)
            ln.bias.copy_(bias)
        else:
            bias = None

    assert ln.bias is None or ln.bias.requires_grad == needs_bias
    assert ln.weight is None or ln.weight.requires_grad == needs_weight

    x = torch.randn(16, 1024, train_config.model.d_model)
    x.requires_grad = False
    y_expected = F.layer_norm(x, [train_config.model.d_model], weight, bias)

    y_actual = ln(x)
    t.assert_allclose(y_actual, y_expected)
//  [python]
def test_layer_norm(train_config: TrainConfig, elementwise_affine: bool, include_bias: bool):
    train_config.model.layer_norm_with_affine = elementwise_affine
    train_config.model.include_bias = include_bias
    ln = LayerNorm.build(train_config.model)

    needs_weight = elementwise_affine
    needs_bias = elementwise_affine and include_bias
    with torch.no_grad():
        if needs_weight:
            weight = torch.randn(train_config.model.d_model)
            ln.weight.copy_(weight)
        else:
            weight = None

        if needs_bias:
            bias = torch.randn(train_config.model.d_model)
            ln.bias.copy_(bias)
        else:
            bias = None

    assert ln.bias is None or ln.bias.requires_grad == needs_bias
    assert ln.weight is None or ln.weight.requires_grad == needs_weight

    x = torch.randn(16, 1024, train_config.model.d_model)
    x.requires_grad = False
    y_expected = F.layer_norm(x, [train_config.model.d_model], weight, bias)

    y_actual = ln(x)
    torch.testing.assert_close(y_actual, y_expected)
//  [python]
def test_maxvit_window_partition(self):
        input_shape = (1, 3, 224, 224)
        partition_size = 7
        n_partitions = input_shape[3] // partition_size

        x = torch.randn(input_shape)

        partition = WindowPartition()
        departition = WindowDepartition()

        x_hat = partition(x, partition_size)
        x_hat = departition(x_hat, partition_size, n_partitions, n_partitions)

        torch.testing.assert_allclose(x, x_hat)
//  [python]
def test_maxvit_window_partition(self):
        input_shape = (1, 3, 224, 224)
        partition_size = 7
        n_partitions = input_shape[3] // partition_size

        x = torch.randn(input_shape)

        partition = WindowPartition()
        departition = WindowDepartition()

        x_hat = partition(x, partition_size)
        x_hat = departition(x_hat, partition_size, n_partitions, n_partitions)

        torch.testing.assert_close(x, x_hat)
//  [python]
def test_matd3_getAction_epsilon_greedy_cnn(
    epsilon, state_dims, action_dims, discrete_actions, device
):
    agent_ids = ["agent_0", "agent_1"]
    net_config = {
        "arch": "cnn",
        "hidden_size": [64, 64],
        "channel_size": [16],
        "kernel_size": [3],
        "stride_size": [1],
        "normalize": False,
    }
    state = {agent: np.random.randn(1, *state_dims[0]) for agent in agent_ids}
    matd3 = MATD3(
        state_dims,
        action_dims,
        one_hot=False,
        n_agents=2,
        agent_ids=agent_ids,
        net_config=net_config,
        max_action=[[1], [1]],
        min_action=[[-1], [0]],
        discrete_actions=discrete_actions,
        device=device,
    )
    cont_actions, discrete_action = matd3.getAction(state, epsilon)
    if discrete_actions:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            assert np.isclose(sum(action), 1.00)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
        for idx, action in enumerate(list(discrete_action.values())):
            assert action <= action_dims[idx] - 1
    else:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
//  [python]
def test_matd3_getAction_epsilon_greedy_cnn(
    epsilon, state_dims, action_dims, discrete_actions, device
):
    agent_ids = ["agent_0", "agent_1"]
    net_config = {
        "arch": "cnn",
        "hidden_size": [64, 64],
        "channel_size": [16],
        "kernel_size": [3],
        "stride_size": [1],
        "normalize": False,
    }
    state = {agent: np.random.randn(1, *state_dims[0]) for agent in agent_ids}
    matd3 = MATD3(
        state_dims,
        action_dims,
        one_hot=False,
        n_agents=2,
        agent_ids=agent_ids,
        net_config=net_config,
        max_action=[[1], [1]],
        min_action=[[-1], [0]],
        discrete_actions=discrete_actions,
        device=device,
    )
    cont_actions, discrete_action = matd3.getAction(state, epsilon)
    if discrete_actions:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            torch.testing.assert_close(sum(action), 1.00)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
        for idx, action in enumerate(list(discrete_action.values())):
            assert action <= action_dims[idx] - 1
    else:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
//  [python]
def test_matd3_getAction_epsilon_greedy_mlp(
    epsilon, state_dims, action_dims, discrete_actions, one_hot, device
):
    agent_ids = ["agent_0", "agent_1"]
    if one_hot:
        state = {
            agent: np.random.randint(0, state_dims[0], *state_dims[0])
            for agent in agent_ids
        }
    else:
        state = {agent: np.random.randn(*state_dims[0]) for agent in agent_ids}
    matd3 = MATD3(
        state_dims,
        action_dims,
        one_hot=one_hot,
        net_config={"arch": "mlp", "hidden_size": [64, 64]},
        n_agents=2,
        agent_ids=agent_ids,
        max_action=[[1], [1]],
        min_action=[[-1], [-1]],
        discrete_actions=discrete_actions,
        device=device,
    )
    cont_actions, discrete_action = matd3.getAction(state, epsilon)
    for idx, action in enumerate(list(cont_actions.values())):
        if one_hot:
            assert len(action[0]) == action_dims[idx]
            assert len(action) == state_dims[idx][0]
            if discrete_actions:
                t.assert_allclose(
                    sum(np.sum(action, 0)), float(state_dims[idx][0])
                )
            act = action[idx]
            assert act.dtype == np.float32
            assert -1 <= act.all() <= 1
        else:
            assert len(action) == action_dims[idx]
            if discrete_actions:
                t.assert_allclose(sum(action), 1.0)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
    if discrete_actions:
        for idx, action in enumerate(list(discrete_action.values())):
            if one_hot:
                assert action.all() <= action_dims[idx] - 1
            else:
                assert action <= action_dims[idx] - 1
//  [python]
def test_matd3_getAction_epsilon_greedy_mlp(
    epsilon, state_dims, action_dims, discrete_actions, one_hot, device
):
    agent_ids = ["agent_0", "agent_1"]
    if one_hot:
        state = {
            agent: np.random.randint(0, state_dims[0], *state_dims[0])
            for agent in agent_ids
        }
    else:
        state = {agent: np.random.randn(*state_dims[0]) for agent in agent_ids}
    matd3 = MATD3(
        state_dims,
        action_dims,
        one_hot=one_hot,
        net_config={"arch": "mlp", "hidden_size": [64, 64]},
        n_agents=2,
        agent_ids=agent_ids,
        max_action=[[1], [1]],
        min_action=[[-1], [-1]],
        discrete_actions=discrete_actions,
        device=device,
    )
    cont_actions, discrete_action = matd3.getAction(state, epsilon)
    for idx, action in enumerate(list(cont_actions.values())):
        if one_hot:
            assert len(action[0]) == action_dims[idx]
            assert len(action) == state_dims[idx][0]
            if discrete_actions:
                torch.testing.assert_close(
                    sum(np.sum(action, 0)), float(state_dims[idx][0])
                )
            act = action[idx]
            assert act.dtype == np.float32
            assert -1 <= act.all() <= 1
        else:
            assert len(action) == action_dims[idx]
            if discrete_actions:
                torch.testing.assert_close(sum(action), 1.0)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
    if discrete_actions:
        for idx, action in enumerate(list(discrete_action.values())):
            if one_hot:
                assert action.all() <= action_dims[idx] - 1
            else:
                assert action <= action_dims[idx] - 1
//  [python]
def test_matd3_getAction_epsilon_greedy_distributed(
    epsilon, state_dims, action_dims, discrete_actions
):
    accelerator = Accelerator()
    agent_ids = ["agent_0", "agent_1"]
    state = {agent: np.random.randn(*state_dims[0]) for agent in agent_ids}
    matd3 = MATD3(
        state_dims,
        action_dims,
        one_hot=False,
        n_agents=2,
        agent_ids=agent_ids,
        max_action=[[1], [1]],
        min_action=[[-1], [-1]],
        discrete_actions=discrete_actions,
        accelerator=accelerator,
    )
    new_actors = [
        DummyEvolvableMLP(
            num_inputs=actor.num_inputs,
            num_outputs=actor.num_outputs,
            hidden_size=actor.hidden_size,
            device=actor.device,
            mlp_output_activation=actor.mlp_output_activation,
        )
        for actor in matd3.actors
    ]
    matd3.actors = new_actors
    cont_actions, discrete_action = matd3.getAction(state, epsilon)
    if discrete_actions:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            assert t.allclose(sum(action), 1.00)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
        for idx, action in enumerate(list(discrete_action.values())):
            assert action <= action_dims[idx] - 1
    else:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
//  [python]
def test_matd3_getAction_epsilon_greedy_distributed(
    epsilon, state_dims, action_dims, discrete_actions
):
    accelerator = Accelerator()
    agent_ids = ["agent_0", "agent_1"]
    state = {agent: np.random.randn(*state_dims[0]) for agent in agent_ids}
    matd3 = MATD3(
        state_dims,
        action_dims,
        one_hot=False,
        n_agents=2,
        agent_ids=agent_ids,
        max_action=[[1], [1]],
        min_action=[[-1], [-1]],
        discrete_actions=discrete_actions,
        accelerator=accelerator,
    )
    new_actors = [
        DummyEvolvableMLP(
            num_inputs=actor.num_inputs,
            num_outputs=actor.num_outputs,
            hidden_size=actor.hidden_size,
            device=actor.device,
            mlp_output_activation=actor.mlp_output_activation,
        )
        for actor in matd3.actors
    ]
    matd3.actors = new_actors
    cont_actions, discrete_action = matd3.getAction(state, epsilon)
    if discrete_actions:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            torch.testing.assert_close(sum(action), 1.00)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
        for idx, action in enumerate(list(discrete_action.values())):
            assert action <= action_dims[idx] - 1
    else:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
//  [python]
def test_maxvit_grid_partition(self):
        input_shape = (1, 3, 224, 224)
        partition_size = 7
        n_partitions = input_shape[3] // partition_size

        x = torch.randn(input_shape)
        pre_swap = SwapAxes(-2, -3)
        post_swap = SwapAxes(-2, -3)

        partition = WindowPartition()
        departition = WindowDepartition()

        x_hat = partition(x, n_partitions)
        x_hat = pre_swap(x_hat)
        x_hat = post_swap(x_hat)
        x_hat = departition(x_hat, n_partitions, partition_size, partition_size)

        torch.testing.assert_allclose(x, x_hat)
//  [python]
def test_maxvit_grid_partition(self):
        input_shape = (1, 3, 224, 224)
        partition_size = 7
        n_partitions = input_shape[3] // partition_size

        x = torch.randn(input_shape)
        pre_swap = SwapAxes(-2, -3)
        post_swap = SwapAxes(-2, -3)

        partition = WindowPartition()
        departition = WindowDepartition()

        x_hat = partition(x, n_partitions)
        x_hat = pre_swap(x_hat)
        x_hat = post_swap(x_hat)
        x_hat = departition(x_hat, n_partitions, partition_size, partition_size)

        torch.testing.assert_close(x, x_hat)
//  [python]
def test_invert_permute(
        self,
        permute_size: int,
        long_index: bool,
    ) -> None:
        index_dtype = torch.int64 if long_index else torch.int32
        permute_list = list(range(permute_size))
        random.shuffle(permute_list)
        inversed_permute_list = [0] * len(permute_list)
        for i in range(permute_size):
            inversed_permute_list[permute_list[i]] = i
        permute = torch.IntTensor(permute_list).type(index_dtype)
        inverse_permute_ref = torch.IntTensor(inversed_permute_list).type(index_dtype)

        inverse_permute_cpu = torch.ops.fbgemm.invert_permute(permute)
        torch.testing.assert_allclose(inverse_permute_cpu, inverse_permute_ref)

        if gpu_available:
            inverse_permute_gpu = torch.ops.fbgemm.invert_permute(permute.cuda())
            torch.testing.assert_allclose(inverse_permute_gpu.cpu(), inverse_permute_cpu)
//  [python]
def test_invert_permute(
        self,
        permute_size: int,
        long_index: bool,
    ) -> None:
        index_dtype = torch.int64 if long_index else torch.int32
        permute_list = list(range(permute_size))
        random.shuffle(permute_list)
        inversed_permute_list = [0] * len(permute_list)
        for i in range(permute_size):
            inversed_permute_list[permute_list[i]] = i
        permute = torch.IntTensor(permute_list).type(index_dtype)
        inverse_permute_ref = torch.IntTensor(inversed_permute_list).type(index_dtype)

        inverse_permute_cpu = torch.ops.fbgemm.invert_permute(permute)
        torch.testing.assert_close(inverse_permute_cpu, inverse_permute_ref)

        if gpu_available:
            inverse_permute_gpu = torch.ops.fbgemm.invert_permute(permute.cuda())
            torch.testing.assert_close(inverse_permute_gpu.cpu(), inverse_permute_cpu)
//  [python]
def test_block_groups():
    model_with_block_groups = OLMo(ModelConfig(d_model=128, n_heads=2, n_layers=9, block_group_size=3)).eval()
    model_without_block_groups = OLMo(ModelConfig(d_model=128, n_heads=2, n_layers=9, block_group_size=1)).eval()

    # We should be able to load the state dict from one model into the other, and vice-versa.
    state_dict_to_load, og_keys_to_new_keys = model_with_block_groups._make_state_dict_compatible(
        model_without_block_groups.state_dict()
    )
    assert og_keys_to_new_keys["transformer.blocks.0.attn_out.weight"] == {
        "transformer.block_groups.0.0.attn_out.weight"
    }
    model_with_block_groups.load_state_dict(state_dict_to_load)

    state_dict_to_load, og_keys_to_new_keys = model_without_block_groups._make_state_dict_compatible(
        model_with_block_groups.state_dict()
    )
    assert og_keys_to_new_keys["transformer.block_groups.0.0.attn_out.weight"] == {
        "transformer.blocks.0.attn_out.weight"
    }
    model_without_block_groups.load_state_dict(state_dict_to_load)

    # Check that output is exactly the same.
    input_ids = torch.randint(0, model_with_block_groups.config.vocab_size, (2, 16))
    with torch.no_grad():
        block_groups_output = model_with_block_groups(input_ids)
        no_block_groups_output = model_without_block_groups(input_ids)

    torch.testing.assert_allclose(block_groups_output, no_block_groups_output)
//  [python]
def test_block_groups():
    model_with_block_groups = OLMo(ModelConfig(d_model=128, n_heads=2, n_layers=9, block_group_size=3)).eval()
    model_without_block_groups = OLMo(ModelConfig(d_model=128, n_heads=2, n_layers=9, block_group_size=1)).eval()

    # We should be able to load the state dict from one model into the other, and vice-versa.
    state_dict_to_load, og_keys_to_new_keys = model_with_block_groups._make_state_dict_compatible(
        model_without_block_groups.state_dict()
    )
    assert og_keys_to_new_keys["transformer.blocks.0.attn_out.weight"] == {
        "transformer.block_groups.0.0.attn_out.weight"
    }
    model_with_block_groups.load_state_dict(state_dict_to_load)

    state_dict_to_load, og_keys_to_new_keys = model_without_block_groups._make_state_dict_compatible(
        model_with_block_groups.state_dict()
    )
    assert og_keys_to_new_keys["transformer.block_groups.0.0.attn_out.weight"] == {
        "transformer.blocks.0.attn_out.weight"
    }
    model_without_block_groups.load_state_dict(state_dict_to_load)

    # Check that output is exactly the same.
    input_ids = torch.randint(0, model_with_block_groups.config.vocab_size, (2, 16))
    with torch.no_grad():
        block_groups_output = model_with_block_groups(input_ids)
        no_block_groups_output = model_without_block_groups(input_ids)

    torch.testing.assert_close(block_groups_output, no_block_groups_output)
//  [python]
def test_segment_sum_csr_empty_input(self) -> None:
        segment_sum_cpu = torch.ops.fbgemm.segment_sum_csr(
            0,
            torch.IntTensor([0]),
            torch.Tensor([]),
        )
        torch.testing.assert_allclose(segment_sum_cpu.numel(), 0, rtol=0, atol=0)

        if torch.cuda.is_available():
            segment_sum_cuda = torch.ops.fbgemm.segment_sum_csr(
                0,
                torch.IntTensor([0]).cuda(),
                torch.Tensor([]).cuda(),
            )
            torch.testing.assert_allclose(
                segment_sum_cuda.cpu().numel(), 0, rtol=0, atol=0
            )
//  [python]
def test_segment_sum_csr_empty_input(self) -> None:
        segment_sum_cpu = torch.ops.fbgemm.segment_sum_csr(
            0,
            torch.IntTensor([0]),
            torch.Tensor([]),
        )
        torch.testing.assert_close(segment_sum_cpu.numel(), 0, rtol=0, atol=0)

        if torch.cuda.is_available():
            segment_sum_cuda = torch.ops.fbgemm.segment_sum_csr(
                0,
                torch.IntTensor([0]).cuda(),
                torch.Tensor([]).cuda(),
            )
            torch.testing.assert_close(
                segment_sum_cuda.cpu().numel(), 0, rtol=0, atol=0
            )
//  [python]
def test_offsets_range(
        self,
        N: int,
        # pyre-fixme[11]: Annotation `int32` is not defined as a type.
        # pyre-fixme[11]: Annotation `int64` is not defined as a type.
        offsets_type: "Union[Type[torch.int32], Type[torch.int64]]",
    ) -> None:
        lengths = np.array([np.random.randint(low=0, high=20) for _ in range(N)])
        offsets = np.cumsum(np.concatenate(([0], lengths)))[:-1]
        range_ref = torch.from_numpy(
            np.concatenate([np.arange(size) for size in lengths])
        )
        output_size = np.sum(lengths)

        offsets_cpu = torch.tensor(offsets, dtype=offsets_type)
        range_cpu = torch.ops.fbgemm.offsets_range(offsets_cpu, output_size)
        range_ref = range_ref.to(range_cpu.dtype)
        torch.testing.assert_allclose(range_cpu, range_ref, rtol=0, atol=0)

        if gpu_available:
            range_gpu = torch.ops.fbgemm.offsets_range(offsets_cpu.cuda(), output_size)
            range_ref = range_ref.to(range_gpu.dtype)
            torch.testing.assert_allclose(range_gpu.cpu(), range_ref, rtol=0, atol=0)
//  [python]
def test_offsets_range(
        self,
        N: int,
        # pyre-fixme[11]: Annotation `int32` is not defined as a type.
        # pyre-fixme[11]: Annotation `int64` is not defined as a type.
        offsets_type: "Union[Type[torch.int32], Type[torch.int64]]",
    ) -> None:
        lengths = np.array([np.random.randint(low=0, high=20) for _ in range(N)])
        offsets = np.cumsum(np.concatenate(([0], lengths)))[:-1]
        range_ref = torch.from_numpy(
            np.concatenate([np.arange(size) for size in lengths])
        )
        output_size = np.sum(lengths)

        offsets_cpu = torch.tensor(offsets, dtype=offsets_type)
        range_cpu = torch.ops.fbgemm.offsets_range(offsets_cpu, output_size)
        range_ref = range_ref.to(range_cpu.dtype)
        torch.testing.assert_close(range_cpu, range_ref, rtol=0, atol=0)

        if gpu_available:
            range_gpu = torch.ops.fbgemm.offsets_range(offsets_cpu.cuda(), output_size)
            range_ref = range_ref.to(range_gpu.dtype)
            torch.testing.assert_close(range_gpu.cpu(), range_ref, rtol=0, atol=0)
//  [python]
def test_segment_sum_csr(self) -> None:
        segment_sum_cpu = torch.ops.fbgemm.segment_sum_csr(
            2,
            torch.IntTensor([0, 2, 3, 5]),
            torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]),
        )
        torch.testing.assert_allclose(
            segment_sum_cpu, torch.Tensor([10.0, 11.0, 34.0]), rtol=0, atol=0
        )
        if torch.cuda.is_available():
            segment_sum_cuda = torch.ops.fbgemm.segment_sum_csr(
                2,
                torch.IntTensor([0, 2, 3, 5]).cuda(),
                torch.Tensor(
                    [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
                ).cuda(),
            )
            torch.testing.assert_allclose(
                segment_sum_cuda.cpu(), torch.Tensor([10.0, 11.0, 34.0]), rtol=0, atol=0
            )
//  [python]
def test_segment_sum_csr(self) -> None:
        segment_sum_cpu = torch.ops.fbgemm.segment_sum_csr(
            2,
            torch.IntTensor([0, 2, 3, 5]),
            torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]),
        )
        torch.testing.assert_close(
            segment_sum_cpu, torch.Tensor([10.0, 11.0, 34.0]), rtol=0, atol=0
        )
        if torch.cuda.is_available():
            segment_sum_cuda = torch.ops.fbgemm.segment_sum_csr(
                2,
                torch.IntTensor([0, 2, 3, 5]).cuda(),
                torch.Tensor(
                    [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
                ).cuda(),
            )
            torch.testing.assert_close(
                segment_sum_cuda.cpu(), torch.Tensor([10.0, 11.0, 34.0]), rtol=0, atol=0
            )
//  [python]
def test_get_tensors_from_safetensors(self) -> None:
        path = Path("example.safetensors")
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }
        save_file(expected_tensors, path)

        tensors = get_tensors(path)

        self.assertIsInstance(tensors, dict)
        self.assertEqual(tensors, expected_tensors)
        path.unlink()
//  [python]
def test_get_tensors_from_safetensors(self) -> None:
        path = Path("example.safetensors")
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }
        save_file(expected_tensors, path)

        tensors = get_tensors(path)

        self.assertIsInstance(tensors, dict)
        torch.testing.assert_close(tensors, expected_tensors)
        path.unlink()
//  [python]
def test_permute102_baddbmm_permute102(
        self,
        batch_size: int,
        m: int,
        k: int,
        n: int,
        use_cpu: bool,
    ) -> None:
        # baddbmm doesn't support half
        dtype = torch.float if use_cpu else torch.half
        device = torch.device("cpu" if use_cpu else "cuda")

        A = torch.rand((m, batch_size, k), dtype=dtype, device=device)
        B = torch.rand((batch_size, k, n), dtype=dtype, device=device)
        # bias_permute102 = torch.rand(batch_size, 1, n).half().cuda()
        # bias = bias_permute102.permute(1, 0, 2)

        bias = torch.rand((batch_size, n), dtype=dtype, device=device)
        bias_permute102 = bias.unsqueeze(1)
        # bias = bias_short.unsqueeze(0)

        A_permute102 = A.permute(1, 0, 2)
        C_permute102 = torch.baddbmm(bias_permute102, A_permute102, B)
        C_ref = C_permute102.permute(1, 0, 2)  # (m, batch_size, n)

        C = torch.ops.fbgemm.permute102_baddbmm_permute102(bias, A, B)
        torch.testing.assert_allclose(C.cpu(), C_ref.cpu())
//  [python]
def test_permute102_baddbmm_permute102(
        self,
        batch_size: int,
        m: int,
        k: int,
        n: int,
        use_cpu: bool,
    ) -> None:
        # baddbmm doesn't support half
        dtype = torch.float if use_cpu else torch.half
        device = torch.device("cpu" if use_cpu else "cuda")

        A = torch.rand((m, batch_size, k), dtype=dtype, device=device)
        B = torch.rand((batch_size, k, n), dtype=dtype, device=device)
        # bias_permute102 = torch.rand(batch_size, 1, n).half().cuda()
        # bias = bias_permute102.permute(1, 0, 2)

        bias = torch.rand((batch_size, n), dtype=dtype, device=device)
        bias_permute102 = bias.unsqueeze(1)
        # bias = bias_short.unsqueeze(0)

        A_permute102 = A.permute(1, 0, 2)
        C_permute102 = torch.baddbmm(bias_permute102, A_permute102, B)
        C_ref = C_permute102.permute(1, 0, 2)  # (m, batch_size, n)

        C = torch.ops.fbgemm.permute102_baddbmm_permute102(bias, A, B)
        torch.testing.assert_close(C.cpu(), C_ref.cpu())
//  [python]
def test_get_tensors_from_pt(self) -> None:
        file_path = "example.pt"
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }
        torch.save(expected_tensors, file_path)

        tensors = get_tensors(file_path)

        self.assertIsInstance(tensors, dict)
        self.assertEqual(tensors, expected_tensors)
        os.remove(file_path)
//  [python]
def test_get_tensors_from_pt(self) -> None:
        file_path = "example.pt"
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }
        torch.save(expected_tensors, file_path)

        tensors = get_tensors(file_path)

        self.assertIsInstance(tensors, dict)
        torch.testing.assert_close(tensors, expected_tensors)
        os.remove(file_path)
//  [python]
def test_sdp_attn_without_xformers_and_pt2_sdp_attn(self) -> None:
        q = torch.randn(3, 32, 64)
        k = torch.randn(3, 32, 64)
        v = torch.randn(3, 32, 64)
        training = True
        mask = torch.randint(0, 2, (3, 32, 32)).bool()
        dropout = 0.0

        with patch(
            "cflearn.toolkit.try_run_xformers_sdp_attn",
            return_value=None,
        ), patch.object(
            F,
            "scaled_dot_product_attention",
            None,
        ):
            output = sdp_attn(q, k, v, training, mask, dropout)

        self.assertIsInstance(output, Tensor)
        self.assertEqual(output.shape, (3, 32, 64))
        expected = F.scaled_dot_product_attention(q, k, v, mask, dropout)
        torch.testing.assert_allclose(output, expected)
//  [python]
def test_sdp_attn_without_xformers_and_pt2_sdp_attn(self) -> None:
        q = torch.randn(3, 32, 64)
        k = torch.randn(3, 32, 64)
        v = torch.randn(3, 32, 64)
        training = True
        mask = torch.randint(0, 2, (3, 32, 32)).bool()
        dropout = 0.0

        with patch(
            "cflearn.toolkit.try_run_xformers_sdp_attn",
            return_value=None,
        ), patch.object(
            F,
            "scaled_dot_product_attention",
            None,
        ):
            output = sdp_attn(q, k, v, training, mask, dropout)

        self.assertIsInstance(output, Tensor)
        self.assertEqual(output.shape, (3, 32, 64))
        expected = F.scaled_dot_product_attention(q, k, v, mask, dropout)
        torch.testing.assert_close(output, expected)
//  [python]
def test_bottom_unique_k_per_row(
        self,
        T: int,
        B: int,
        L: int,
    ) -> None:
        E = 1000000
        all_indices = (np.random.zipf(a=1.15, size=(T, B, 3 * L)) - 1) % E
        all_indices_deduped = torch.ops.fbgemm.bottom_k_per_row(
            torch.as_tensor(all_indices), torch.tensor([0, L], dtype=torch.long), True
        )
        for index_tuple in itertools.product(range(T), range(B)):
            # sample without replacement from
            # https://stats.stackexchange.com/questions/20590/how-do-i-sample-without-replacement-using-a-sampling-with-replacement-function
            r = set()
            for x in all_indices[index_tuple]:
                if x not in r:
                    r.add(x)
                    if len(r) == L:
                        break
            assert (len(r)) == L, "too skewed distribution (alpha too big)"
            all_indices[index_tuple][:L] = sorted(r)
        all_indices_deduped_ref = torch.as_tensor(all_indices[:, :, :L])
        torch.testing.assert_allclose(all_indices_deduped, all_indices_deduped_ref)
//  [python]
def test_bottom_unique_k_per_row(
        self,
        T: int,
        B: int,
        L: int,
    ) -> None:
        E = 1000000
        all_indices = (np.random.zipf(a=1.15, size=(T, B, 3 * L)) - 1) % E
        all_indices_deduped = torch.ops.fbgemm.bottom_k_per_row(
            torch.as_tensor(all_indices), torch.tensor([0, L], dtype=torch.long), True
        )
        for index_tuple in itertools.product(range(T), range(B)):
            # sample without replacement from
            # https://stats.stackexchange.com/questions/20590/how-do-i-sample-without-replacement-using-a-sampling-with-replacement-function
            r = set()
            for x in all_indices[index_tuple]:
                if x not in r:
                    r.add(x)
                    if len(r) == L:
                        break
            assert (len(r)) == L, "too skewed distribution (alpha too big)"
            all_indices[index_tuple][:L] = sorted(r)
        all_indices_deduped_ref = torch.as_tensor(all_indices[:, :, :L])
        torch.testing.assert_close(all_indices_deduped, all_indices_deduped_ref)
//  [python]
def test_get_tensors_from_state_dict(self) -> None:
        state_dict = {
            "state_dict": {
                "tensor1": torch.tensor([1, 2, 3]),
                "tensor2": torch.tensor([4, 5, 6]),
            }
        }
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }

        tensors = get_tensors(state_dict)

        self.assertIsInstance(tensors, dict)
        self.assertEqual(tensors, expected_tensors)
//  [python]
def test_get_tensors_from_state_dict(self) -> None:
        state_dict = {
            "state_dict": {
                "tensor1": torch.tensor([1, 2, 3]),
                "tensor2": torch.tensor([4, 5, 6]),
            }
        }
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }

        tensors = get_tensors(state_dict)

        self.assertIsInstance(tensors, dict)
        torch.testing.assert_close(tensors, expected_tensors)
//  [python]
def test_get_tensors_from_dict(self) -> None:
        tensor_dict = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }

        tensors = get_tensors(tensor_dict)

        self.assertIsInstance(tensors, dict)
        self.assertEqual(tensors, expected_tensors)
//  [python]
def test_get_tensors_from_dict(self) -> None:
        tensor_dict = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }
        expected_tensors = {
            "tensor1": torch.tensor([1, 2, 3]),
            "tensor2": torch.tensor([4, 5, 6]),
        }

        tensors = get_tensors(tensor_dict)

        self.assertIsInstance(tensors, dict)
        torch.testing.assert_close(tensors, expected_tensors)
//  [python]
def test_initialize_custom_method(self) -> None:
        initializer = Initializer()
        m = nn.Linear(10, 2)
        param = m.weight

        @Initializer.register("custom")
        def custom(self: Initializer, param: nn.Parameter) -> None:
            with torch.no_grad():
                param.data.fill_(1.0)

        initializer.initialize(param, "custom")

        self.assertIsInstance(param, Tensor)
        t.assert_close(param, torch.ones_like(param))
//  [python]
def test_initialize_custom_method(self) -> None:
        initializer = Initializer()
        m = nn.Linear(10, 2)
        param = m.weight

        @Initializer.register("custom")
        def custom(self: Initializer, param: nn.Parameter) -> None:
            with torch.no_grad():
                param.data.fill_(1.0)

        initializer.initialize(param, "custom")

        self.assertIsInstance(param, Tensor)
        torch.testing.assert_close(param, torch.ones_like(param))
//  [python]
def test_fix_denormal_states(self) -> None:
        states = {
            "a": torch.tensor([1.0, 2.0, 1.0e-33]),
            "b": torch.tensor([4.0, 5.0, 6.0]),
        }
        expected_states = {
            "a": torch.tensor([1.0, 2.0, 0.0]),
            "b": torch.tensor([4.0, 5.0, 6.0]),
        }

        new_states = fix_denormal_states(states)

        self.assertIsInstance(new_states, dict)
        self.assertEqual(new_states, expected_states)
//  [python]
def test_fix_denormal_states(self) -> None:
        states = {
            "a": torch.tensor([1.0, 2.0, 1.0e-33]),
            "b": torch.tensor([4.0, 5.0, 6.0]),
        }
        expected_states = {
            "a": torch.tensor([1.0, 2.0, 0.0]),
            "b": torch.tensor([4.0, 5.0, 6.0]),
        }

        new_states = fix_denormal_states(states)

        self.assertIsInstance(new_states, dict)
        torch.testing.assert_close(new_states, expected_states)
//  [python]
def test_bucketize_sparse_features(
        self,
        index_type: Type[torch.dtype],
        has_weight: bool,
        bucketize_pos: bool,
    ) -> None:
        # pyre-ignore [6]
        lengths = torch.tensor([0, 2, 1, 3], dtype=index_type)
        # pyre-ignore [6]
        indices = torch.tensor([10, 10, 15, 20, 25, 30], dtype=index_type)
        weights = (
            torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=torch.float)
            if has_weight
            else None
        )

        # pyre-ignore [6]
        new_lengths_ref = torch.tensor([0, 2, 0, 2, 0, 0, 1, 1], dtype=index_type)
        # pyre-ignore [6]
        new_indices_ref = torch.tensor([5, 5, 10, 15, 7, 12], dtype=index_type)
        new_weights_ref = torch.tensor(
            [1.0, 2.0, 4.0, 6.0, 3.0, 5.0], dtype=torch.float
        )
        # pyre-ignore [6]
        new_pos_ref = torch.tensor([0, 1, 0, 2, 0, 1], dtype=index_type)
        (
            new_lengths_cpu,
            new_indices_cpu,
            new_weights_cpu,
            new_pos_cpu,
        ) = torch.ops.fbgemm.bucketize_sparse_features(
            lengths, indices, bucketize_pos, 2, weights
        )
        t.assert_close(new_lengths_cpu, new_lengths_ref, rtol=0, atol=0)
        t.assert_close(new_indices_cpu, new_indices_ref, rtol=0, atol=0)
        if has_weight:
            t.assert_close(new_weights_cpu, new_weights_ref)
        if bucketize_pos:
            t.assert_close(new_pos_cpu, new_pos_ref)
        if gpu_available:
            (
                new_lengths_gpu,
                new_indices_gpu,
                new_weights_gpu,
                new_pos_gpu,
            ) = torch.ops.fbgemm.bucketize_sparse_features(
                lengths.cuda(),
                indices.cuda(),
                bucketize_pos,
                2,
                # pyre-fixme[16]: `Optional` has no attribute `cuda`.
                weights.cuda() if has_weight else None,
            )
            t.assert_close(
                new_lengths_gpu.cpu(), new_lengths_ref, rtol=0, atol=0
            )
            t.assert_close(
                new_indices_gpu.cpu(), new_indices_ref, rtol=0, atol=0
            )
            if has_weight:
                t.assert_close(new_weights_gpu.cpu(), new_weights_cpu)
            if bucketize_pos:
                t.assert_close(new_pos_gpu.cpu(), new_pos_cpu)
//  [python]
def test_bucketize_sparse_features(
        self,
        index_type: Type[torch.dtype],
        has_weight: bool,
        bucketize_pos: bool,
    ) -> None:
        # pyre-ignore [6]
        lengths = torch.tensor([0, 2, 1, 3], dtype=index_type)
        # pyre-ignore [6]
        indices = torch.tensor([10, 10, 15, 20, 25, 30], dtype=index_type)
        weights = (
            torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=torch.float)
            if has_weight
            else None
        )

        # pyre-ignore [6]
        new_lengths_ref = torch.tensor([0, 2, 0, 2, 0, 0, 1, 1], dtype=index_type)
        # pyre-ignore [6]
        new_indices_ref = torch.tensor([5, 5, 10, 15, 7, 12], dtype=index_type)
        new_weights_ref = torch.tensor(
            [1.0, 2.0, 4.0, 6.0, 3.0, 5.0], dtype=torch.float
        )
        # pyre-ignore [6]
        new_pos_ref = torch.tensor([0, 1, 0, 2, 0, 1], dtype=index_type)
        (
            new_lengths_cpu,
            new_indices_cpu,
            new_weights_cpu,
            new_pos_cpu,
        ) = torch.ops.fbgemm.bucketize_sparse_features(
            lengths, indices, bucketize_pos, 2, weights
        )
        torch.testing.assert_close(new_lengths_cpu, new_lengths_ref, rtol=0, atol=0)
        torch.testing.assert_close(new_indices_cpu, new_indices_ref, rtol=0, atol=0)
        if has_weight:
            torch.testing.assert_close(new_weights_cpu, new_weights_ref)
        if bucketize_pos:
            torch.testing.assert_close(new_pos_cpu, new_pos_ref)
        if gpu_available:
            (
                new_lengths_gpu,
                new_indices_gpu,
                new_weights_gpu,
                new_pos_gpu,
            ) = torch.ops.fbgemm.bucketize_sparse_features(
                lengths.cuda(),
                indices.cuda(),
                bucketize_pos,
                2,
                # pyre-fixme[16]: `Optional` has no attribute `cuda`.
                weights.cuda() if has_weight else None,
            )
            torch.testing.assert_close(
                new_lengths_gpu.cpu(), new_lengths_ref, rtol=0, atol=0
            )
            torch.testing.assert_close(
                new_indices_gpu.cpu(), new_indices_ref, rtol=0, atol=0
            )
            if has_weight:
                torch.testing.assert_close(new_weights_gpu.cpu(), new_weights_cpu)
            if bucketize_pos:
                torch.testing.assert_close(new_pos_gpu.cpu(), new_pos_cpu)
//  [python]
def test_slerp(self) -> None:
        x1 = torch.tensor([[1.0, 0.0]])
        x2 = torch.tensor([[0.0, 1.0]])
        r1 = 0.5

        result = slerp(x1, x2, r1)

        self.assertIsInstance(result, torch.Tensor)
        self.assertEqual(result.shape, (1, 2))
        self.assertTrue(torch.allclose(result, torch.tensor([[0.7071, 0.7071]])))
//  [python]
def test_slerp(self) -> None:
        x1 = torch.tensor([[1.0, 0.0]])
        x2 = torch.tensor([[0.0, 1.0]])
        r1 = 0.5

        result = slerp(x1, x2, r1)

        self.assertIsInstance(result, torch.Tensor)
        self.assertEqual(result.shape, (1, 2))
        torch.testing.assert_close(result, torch.tensor([[0.7071, 0.7071]]))
//  [python]
def test_mean_std(self) -> None:
        latent_map = torch.rand(1, 3, 4, 4)

        mean, std = mean_std(latent_map)
        expected_mean = latent_map.mean(dim=[0, 2, 3], keepdim=True)
        expected_std = latent_map.std(dim=[0, 2, 3], keepdim=True)

        self.assertIsInstance(mean, Tensor)
        self.assertIsInstance(std, Tensor)
        self.assertTrue(torch.allclose(mean, expected_mean, atol=1.0e-4, rtol=1.0e-4))
        self.assertTrue(torch.allclose(std, expected_std, atol=1.0e-4, rtol=1.0e-4))
//  [python]
def test_mean_std(self) -> None:
        latent_map = torch.rand(1, 3, 4, 4)

        mean, std = mean_std(latent_map)
        expected_mean = latent_map.mean(dim=[0, 2, 3], keepdim=True)
        expected_std = latent_map.std(dim=[0, 2, 3], keepdim=True)

        self.assertIsInstance(mean, Tensor)
        self.assertIsInstance(std, Tensor)
        torch.testing.assert_close(mean, expected_mean, atol=1.0e-4, rtol=1.0e-4)
        torch.testing.assert_close(std, expected_std, atol=1.0e-4, rtol=1.0e-4)
//  [python]
def snippet_torch_consistency(op, torch_op, sample):
    thunder_result = op(*sample.args, **sample.kwargs)
    torch_result = torch_op(*sample.args, **sample.kwargs)
    # TODO: Pass a custom_comparator which has higher tol for bf16 cases
    t.assert_close(thunder_result, torch_result, equal_nan=True, atol=0.0625, rtol=5e-2)
//  [python]
def snippet_torch_consistency(op, torch_op, sample):
    thunder_result = op(*sample.args, **sample.kwargs)
    torch_result = torch_op(*sample.args, **sample.kwargs)
    # TODO: Pass a custom_comparator which has higher tol for bf16 cases
    assert_close(thunder_result, torch_result, equal_nan=True, atol=0.0625, rtol=5e-2)
//  [python]
def assert_tensor_equal(*args, **kwargs):  # noqa: N802
        t.assert_allclose(
            *args,
            rtol=0,
            atol=0,
            check_dtype=False,
            **kwargs,
        )
//  [python]
def assert_tensor_equal(*args, **kwargs):  # noqa: N802
        torch.testing.assert_close(
            *args,
            rtol=0,
            atol=0,
            check_dtype=False,
            **kwargs,
        )
//  [python]
def _test_throughput_with_input(
        self,
        num_processed: int,
        elapsed_time_sec: float,
    ) -> None:
        t.assert_close(
            throughput(num_processed, elapsed_time_sec),
            torch.tensor(num_processed / elapsed_time_sec),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_throughput_with_input(
        self,
        num_processed: int,
        elapsed_time_sec: float,
    ) -> None:
        torch.testing.assert_close(
            throughput(num_processed, elapsed_time_sec),
            torch.tensor(num_processed / elapsed_time_sec),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def assert_tensor_almost_equal(*args, **kwargs):  # noqa: N802
        t.assert_allclose(
            *args,
            **kwargs,
            check_dtype=False,
        )
//  [python]
def assert_tensor_almost_equal(*args, **kwargs):  # noqa: N802
        torch.testing.assert_close(
            *args,
            **kwargs,
            check_dtype=False,
        )
//  [python]
def test_small_model_logits(self):
        model_id = "hf-internal-testing/Mixtral-tiny"
        dummy_input = torch.LongTensor([[0, 1, 0], [0, 1, 0]]).to(torch_device)

        model = MixtralForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).to(
            torch_device
        )
        # TODO: might need to tweak it in case the logits do not match on our daily runners
        # these logits have been obtained with the original megablocks impelmentation.
        EXPECTED_LOGITS = {
            7: torch.Tensor([[0.1670, 0.1620, 0.6094], [-0.8906, -0.1588, -0.6060], [0.1572, 0.1290, 0.7246]]).to(
                torch_device
            ),
            8: torch.Tensor([[0.1631, 0.1621, 0.6094], [-0.8906, -0.1621, -0.6094], [0.1572, 0.1270, 0.7227]]).to(
                torch_device
            ),
        }
        with torch.no_grad():
            logits = model(dummy_input).logits

        torch.allclose(
            logits[0, :3, :3], EXPECTED_LOGITS[self.cuda_compute_capability_major_version], atol=1e-3, rtol=1e-3
        )
        torch.allclose(
            logits[1, :3, :3], EXPECTED_LOGITS[self.cuda_compute_capability_major_version], atol=1e-3, rtol=1e-3
        )
//  [python]
def test_small_model_logits(self):
        model_id = "hf-internal-testing/Mixtral-tiny"
        dummy_input = torch.LongTensor([[0, 1, 0], [0, 1, 0]]).to(torch_device)

        model = MixtralForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).to(
            torch_device
        )
        # TODO: might need to tweak it in case the logits do not match on our daily runners
        # these logits have been obtained with the original megablocks impelmentation.
        EXPECTED_LOGITS = {
            7: torch.Tensor([[0.1670, 0.1620, 0.6094], [-0.8906, -0.1588, -0.6060], [0.1572, 0.1290, 0.7246]]).to(
                torch_device
            ),
            8: torch.Tensor([[0.1631, 0.1621, 0.6094], [-0.8906, -0.1621, -0.6094], [0.1572, 0.1270, 0.7227]]).to(
                torch_device
            ),
        }
        with torch.no_grad():
            logits = model(dummy_input).logits

        torch.testing.assert_close(
            logits[0, :3, :3], EXPECTED_LOGITS[self.cuda_compute_capability_major_version], atol=1e-3, rtol=1e-3
        )
        torch.testing.assert_close(
            logits[1, :3, :3], EXPECTED_LOGITS[self.cuda_compute_capability_major_version], atol=1e-3, rtol=1e-3
        )
//  [python]
def test_torch_slice(self):
        A = torch.randn((10, 5))
        tensors = {
            "a": A,
        }
        save_file_pt(tensors, "./slice.safetensors")

        # Now loading
        with safe_open("./slice.safetensors", framework="pt", device="cpu") as f:
            slice_ = f.get_slice("a")
            tensor = slice_[:]
            self.assertEqual(list(tensor.shape), [10, 5])
            self.assertTrue(torch.allclose(tensor, A))

            tensor = slice_[:2]
            self.assertEqual(list(tensor.shape), [2, 5])
            self.assertTrue(torch.allclose(tensor, A[:2]))

            tensor = slice_[:, :2]
            self.assertEqual(list(tensor.shape), [10, 2])
            self.assertTrue(torch.allclose(tensor, A[:, :2]))

            tensor = slice_[0, :2]
            self.assertEqual(list(tensor.shape), [2])
            self.assertTrue(torch.allclose(tensor, A[0, :2]))

            tensor = slice_[2:, 0]
            self.assertEqual(list(tensor.shape), [8])
            self.assertTrue(torch.allclose(tensor, A[2:, 0]))

            tensor = slice_[2:, 1]
            self.assertEqual(list(tensor.shape), [8])
            self.assertTrue(torch.allclose(tensor, A[2:, 1]))

            tensor = slice_[2:, -1]
            self.assertEqual(list(tensor.shape), [8])
            self.assertTrue(torch.allclose(tensor, A[2:, -1]))