jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check_parity_ddp(torch_bn, fs_bn, x):
    yh = torch.randn_like(x)
    rank = dist.get_rank()
    torch_ddp = DDP(torch_bn, device_ids=[rank])
    torch_bn = torch_ddp.module
    torch_x = x.detach()
    torch_x.requires_grad = True
    torch_y = torch_ddp(torch_x)
    torch_y.backward(yh)
    fs_ddp = DDP(fs_bn, device_ids=[rank])
    fs_bn = fs_ddp.module
    fs_x = x.detach()
    fs_x.requires_grad = True
    fs_y = fs_ddp(fs_x)
    fs_y.backward(yh)
    torch.testing.assert_close(torch_y, fs_y)
    torch.testing.assert_close(torch_x.grad, fs_x.grad)
    if isinstance(torch_bn, nn.Sequential):
        torch_bn = torch_bn[0]
        fs_bn = fs_bn[0]
    torch.testing.assert_close(torch_bn.running_mean, fs_bn.running_mean)
    torch.testing.assert_close(torch_bn.running_var, fs_bn.running_var)
    torch.testing.assert_close(torch_bn.weight, fs_bn.weight)
    torch.testing.assert_close(torch_bn.bias, fs_bn.bias)
    torch.testing.assert_close(torch_bn.weight.grad, fs_bn.weight.grad)
    torch.testing.assert_close(torch_bn.bias.grad, fs_bn.bias.grad)
//  [python]
def test_ensure_directed_line_graph_compat(graph_data, request):
    s, g, state = request.getfixturevalue(graph_data)
    lattice = (
        torch.tensor(s.lattice.matrix, dtype=matgl.float_th).unsqueeze(dim=0)
        if graph_data != "graph_CH4"
        else torch.tensor(np.identity(3), dtype=matgl.float_th).unsqueeze(dim=0)
    )
    g.edata["pbc_offshift"] = torch.matmul(g.edata["pbc_offset"], lattice[0])
    g.ndata["pos"] = g.ndata["frac_coords"] @ lattice[0]
    bv, bd = compute_pair_vector_and_distance(g)
    g.edata["bond_vec"] = bv
    g.edata["bond_dist"] = bd
    line_graph = create_line_graph(g, 3.0, directed=True)
    edge_ids = line_graph.ndata["edge_ids"].clone()
    src_bond_sign = line_graph.ndata["src_bond_sign"].clone()
    line_graph.ndata["edge_ids"] = torch.zeros(line_graph.num_nodes(), dtype=torch.long)
    line_graph.ndata["src_bond_sign"] = torch.zeros(line_graph.num_nodes())

    assert not torch.allclose(line_graph.ndata["edge_ids"], edge_ids)
    assert not torch.allclose(line_graph.ndata["src_bond_sign"], src_bond_sign)

    # test that the line graph is not compatible
    line_graph = ensure_line_graph_compatibility(g, line_graph, 3.0, directed=True)
    tt.assert_allclose(line_graph.ndata["edge_ids"], edge_ids)
    tt.assert_allclose(line_graph.ndata["src_bond_sign"], src_bond_sign)

    with pytest.raises(RuntimeError):
        ensure_line_graph_compatibility(g, line_graph, 1.0, directed=True)
//  [python]
def test_ensure_directed_line_graph_compat(graph_data, request):
    s, g, state = request.getfixturevalue(graph_data)
    lattice = (
        torch.tensor(s.lattice.matrix, dtype=matgl.float_th).unsqueeze(dim=0)
        if graph_data != "graph_CH4"
        else torch.tensor(np.identity(3), dtype=matgl.float_th).unsqueeze(dim=0)
    )
    g.edata["pbc_offshift"] = torch.matmul(g.edata["pbc_offset"], lattice[0])
    g.ndata["pos"] = g.ndata["frac_coords"] @ lattice[0]
    bv, bd = compute_pair_vector_and_distance(g)
    g.edata["bond_vec"] = bv
    g.edata["bond_dist"] = bd
    line_graph = create_line_graph(g, 3.0, directed=True)
    edge_ids = line_graph.ndata["edge_ids"].clone()
    src_bond_sign = line_graph.ndata["src_bond_sign"].clone()
    line_graph.ndata["edge_ids"] = torch.zeros(line_graph.num_nodes(), dtype=torch.long)
    line_graph.ndata["src_bond_sign"] = torch.zeros(line_graph.num_nodes())

    assert not torch.allclose(line_graph.ndata["edge_ids"], edge_ids)
    assert not torch.allclose(line_graph.ndata["src_bond_sign"], src_bond_sign)

    # test that the line graph is not compatible
    line_graph = ensure_line_graph_compatibility(g, line_graph, 3.0, directed=True)
    tt.assert_close(line_graph.ndata["edge_ids"], edge_ids)
    tt.assert_close(line_graph.ndata["src_bond_sign"], src_bond_sign)

    with pytest.raises(RuntimeError):
        ensure_line_graph_compatibility(g, line_graph, 1.0, directed=True)
//  [python]
def test_collate_tuples(graphs_nx, features_shapes, device):
    graphs_in = [add_random_features(Graph.from_networkx(g), **features_shapes).to(device) for g in graphs_nx]
    graphs_out = list(reversed(graphs_in))
    xs = torch.rand(len(graphs_in), 10, 32)
    ys = torch.rand(len(graphs_in), 7)

    samples = list(zip(graphs_in, xs, ys, graphs_out))
    batch = GraphBatch.collate(samples)

    for g1, g2 in zip(graphs_in, batch[0]):
        assert_graphs_equal(g1, g2)

    torch.testing.assert_allclose(xs, batch[1])

    torch.testing.assert_allclose(ys, batch[2])

    for g1, g2 in zip(graphs_out, batch[3]):
        assert_graphs_equal(g1, g2)
//  [python]
def test_collate_tuples(graphs_nx, features_shapes, device):
    graphs_in = [add_random_features(Graph.from_networkx(g), **features_shapes).to(device) for g in graphs_nx]
    graphs_out = list(reversed(graphs_in))
    xs = torch.rand(len(graphs_in), 10, 32)
    ys = torch.rand(len(graphs_in), 7)

    samples = list(zip(graphs_in, xs, ys, graphs_out))
    batch = GraphBatch.collate(samples)

    for g1, g2 in zip(graphs_in, batch[0]):
        assert_graphs_equal(g1, g2)

    torch.testing.assert_close(xs, batch[1])

    torch.testing.assert_close(ys, batch[2])

    for g1, g2 in zip(graphs_out, batch[3]):
        assert_graphs_equal(g1, g2)
//  [python]
def testGet(self, _, model_builder):
        model = model_builder(self)
        params = higher.utils.get_func_params(model)
        fmodel = higher.patch.monkeypatch(model)
        zipped = zip(model.parameters(), params, fmodel.parameters())
        for true_param, got_param, got_param2 in zipped:
            torch.testing.assert_allclose(true_param, got_param)
//  [python]
def testGet(self, _, model_builder):
        model = model_builder(self)
        params = higher.utils.get_func_params(model)
        fmodel = higher.patch.monkeypatch(model)
        zipped = zip(model.parameters(), params, fmodel.parameters())
        for true_param, got_param, got_param2 in zipped:
            torch.testing.assert_close(true_param, got_param)
//  [python]
def test_collate_dicts(graphs_nx, features_shapes, device):
    graphs_in = [add_random_features(Graph.from_networkx(g), **features_shapes).to(device) for g in graphs_nx]
    graphs_out = list(reversed(graphs_in))
    xs = torch.rand(len(graphs_in), 10, 32)
    ys = torch.rand(len(graphs_in), 7)

    samples = [{'in': gi, 'x': x, 'y': y, 'out': go} for gi, x, y, go in zip(graphs_in, xs, ys, graphs_out)]
    batch = GraphBatch.collate(samples)

    for g1, g2 in zip(graphs_in, batch['in']):
        assert_graphs_equal(g1, g2)

    torch.testing.assert_allclose(xs, batch['x'])

    torch.testing.assert_allclose(ys, batch['y'])

    for g1, g2 in zip(graphs_out, batch['out']):
        assert_graphs_equal(g1, g2)
//  [python]
def test_collate_dicts(graphs_nx, features_shapes, device):
    graphs_in = [add_random_features(Graph.from_networkx(g), **features_shapes).to(device) for g in graphs_nx]
    graphs_out = list(reversed(graphs_in))
    xs = torch.rand(len(graphs_in), 10, 32)
    ys = torch.rand(len(graphs_in), 7)

    samples = [{'in': gi, 'x': x, 'y': y, 'out': go} for gi, x, y, go in zip(graphs_in, xs, ys, graphs_out)]
    batch = GraphBatch.collate(samples)

    for g1, g2 in zip(graphs_in, batch['in']):
        assert_graphs_equal(g1, g2)

    torch.testing.assert_close(xs, batch['x'])

    torch.testing.assert_close(ys, batch['y'])

    for g1, g2 in zip(graphs_out, batch['out']):
        assert_graphs_equal(g1, g2)
//  [python]
def testForward(self, _, model_builder):
        model = model_builder(self)
        fmodel = higher.patch.monkeypatch(model)
        for _ in range(10):
            inputs = torch.rand(8, 4)
            ref_outputs = model(inputs)
            outputs = fmodel(inputs)
            torch.testing.assert_allclose(outputs, ref_outputs)
//  [python]
def testForward(self, _, model_builder):
        model = model_builder(self)
        fmodel = higher.patch.monkeypatch(model)
        for _ in range(10):
            inputs = torch.rand(8, 4)
            ref_outputs = model(inputs)
            outputs = fmodel(inputs)
            torch.testing.assert_close(outputs, ref_outputs)
//  [python]
def test_masking(self):
        m = ConvBlock(1, 1, activation=None, use_masks=True)
        # Reset parameters so that the operation does nothing
        for name, param in m.named_parameters():
            param.requires_grad = False
            param.data.zero_()
            if name == "conv.weight":
                param[:, :, 1, 1] = 1

        x = torch.randn(3, 1, 11, 13)
        y = m(PaddedTensor(x, torch.tensor([[11, 13], [10, 12], [3, 2]]))).data

        # Check sample 1
        torch.testing.assert_allclose(x[0, :, :, :], y[0, :, :, :])
        # Check sample 2
        torch.testing.assert_allclose(x[1, :, :10, :12], y[1, :, :10, :12])
        torch.testing.assert_allclose(torch.zeros(1, 1, 13), y[1, :, 10:, :])
        torch.testing.assert_allclose(torch.zeros(1, 11, 1), y[1, :, :, 12:])
        # Check sample 3
        torch.testing.assert_allclose(x[2, :, :3, :2], y[2, :, :3, :2])
        torch.testing.assert_allclose(torch.zeros(1, 8, 13), y[2, :, 3:, :])
        torch.testing.assert_allclose(torch.zeros(1, 11, 11), y[2, :, :, 2:])
//  [python]
def test_masking(self):
        m = ConvBlock(1, 1, activation=None, use_masks=True)
        # Reset parameters so that the operation does nothing
        for name, param in m.named_parameters():
            param.requires_grad = False
            param.data.zero_()
            if name == "conv.weight":
                param[:, :, 1, 1] = 1

        x = torch.randn(3, 1, 11, 13)
        y = m(PaddedTensor(x, torch.tensor([[11, 13], [10, 12], [3, 2]]))).data

        # Check sample 1
        torch.testing.assert_close(x[0, :, :, :], y[0, :, :, :])
        # Check sample 2
        torch.testing.assert_close(x[1, :, :10, :12], y[1, :, :10, :12])
        torch.testing.assert_close(torch.zeros(1, 1, 13), y[1, :, 10:, :])
        torch.testing.assert_close(torch.zeros(1, 11, 1), y[1, :, :, 12:])
        # Check sample 3
        torch.testing.assert_close(x[2, :, :3, :2], y[2, :, :3, :2])
        torch.testing.assert_close(torch.zeros(1, 8, 13), y[2, :, 3:, :])
        torch.testing.assert_close(torch.zeros(1, 11, 11), y[2, :, :, 2:])
//  [python]
def testBackward(self, _, model_builder):
        model = model_builder(self)
        fmodel = higher.patch.monkeypatch(model)
        for _ in range(10):
            inputs = torch.rand(8, 4)
            ref_loss = model(inputs).sum().pow(2)
            loss = fmodel(inputs).sum().pow(2)
            ref_loss.backward()
            loss.backward()
            for true_param, got_param in zip(
                model.parameters(), fmodel.parameters()
            ):
                if true_param.grad is None:
                    self.assertIsNone(got_param.grad)
                else:
                    self.assertIsNotNone(got_param.grad)
                    torch.testing.assert_allclose(
                        true_param.grad, got_param.grad
                    )
        model.zero_grad()
//  [python]
def testBackward(self, _, model_builder):
        model = model_builder(self)
        fmodel = higher.patch.monkeypatch(model)
        for _ in range(10):
            inputs = torch.rand(8, 4)
            ref_loss = model(inputs).sum().pow(2)
            loss = fmodel(inputs).sum().pow(2)
            ref_loss.backward()
            loss.backward()
            for true_param, got_param in zip(
                model.parameters(), fmodel.parameters()
            ):
                if true_param.grad is None:
                    self.assertIsNone(got_param.grad)
                else:
                    self.assertIsNotNone(got_param.grad)
                    torch.testing.assert_close(
                        true_param.grad, got_param.grad
                    )
        model.zero_grad()
//  [python]
def testRNNForward(self, _, rnn_constructor):
        num_layers = 2
        hidden_size = 20
        num_feats = 10

        batch_size = 3
        seq_length = 5

        for _ in range(10):

            rnn = rnn_constructor(num_feats, hidden_size, num_layers)
            frnn = higher.patch.monkeypatch(rnn)

            inputs = torch.randn(seq_length, batch_size, num_feats)

            output, state = rnn(inputs)
            foutput, fstate = frnn(inputs)

            torch.testing.assert_allclose(output, foutput)
            if isinstance(state, tuple):
                self.assertEqual(len(state), len(fstate))
                for s, fs in zip(state, fstate):
                    torch.testing.assert_allclose(s, fs)
            else:
                torch.testing.assert_allclose(state, fstate)
//  [python]
def testRNNForward(self, _, rnn_constructor):
        num_layers = 2
        hidden_size = 20
        num_feats = 10

        batch_size = 3
        seq_length = 5

        for _ in range(10):

            rnn = rnn_constructor(num_feats, hidden_size, num_layers)
            frnn = higher.patch.monkeypatch(rnn)

            inputs = torch.randn(seq_length, batch_size, num_feats)

            output, state = rnn(inputs)
            foutput, fstate = frnn(inputs)

            torch.testing.assert_close(output, foutput)
            if isinstance(state, tuple):
                self.assertEqual(len(state), len(fstate))
                for s, fs in zip(state, fstate):
                    torch.testing.assert_close(s, fs)
            else:
                torch.testing.assert_close(state, fstate)
//  [python]
def test_write_img(img_: Tensor, ext: str) -> None:
    """Unit test for writing images."""
    with NamedTemporaryFile("w") as f:
        path = f"{f.name}{ext}"
        write_img(img_, path)
        img = read_image(path)
        torch.testing.assert_allclose(img, img_)
//  [python]
def test_write_img(img_: Tensor, ext: str) -> None:
    """Unit test for writing images."""
    with NamedTemporaryFile("w") as f:
        path = f"{f.name}{ext}"
        write_img(img_, path)
        img = read_image(path)
        torch.testing.assert_close(img, img_)
//  [python]
def test_read_img(
    img_: Tensor, ext: str, write_image: Callable[..., Any]
) -> None:
    """Unit test for reading images."""
    with NamedTemporaryFile("w") as f:
        path = f.name + ext
        write_image(img_, path)
        img = read_img(path)
        torch.testing.assert_allclose(img, img_)
//  [python]
def test_read_img(
    img_: Tensor, ext: str, write_image: Callable[..., Any]
) -> None:
    """Unit test for reading images."""
    with NamedTemporaryFile("w") as f:
        path = f.name + ext
        write_image(img_, path)
        img = read_img(path)
        torch.testing.assert_close(img, img_)
//  [python]
def test_get_topk_from_heatmap():
    from mmdet.models.utils.gaussian_target import get_topk_from_heatmap

    from mmdeploy.codebase.mmdet.models.utils.gaussian_target import \
        get_topk_from_heatmap__default
    scores = torch.rand(1, 2, 4, 4)

    gts = get_topk_from_heatmap(scores, k=20)
    outs = get_topk_from_heatmap__default(scores, k=20)

    for gt, out in zip(gts, outs):
        torch.testing.assert_allclose(gt, out)
//  [python]
def test_get_topk_from_heatmap():
    from mmdet.models.utils.gaussian_target import get_topk_from_heatmap

    from mmdeploy.codebase.mmdet.models.utils.gaussian_target import \
        get_topk_from_heatmap__default
    scores = torch.rand(1, 2, 4, 4)

    gts = get_topk_from_heatmap(scores, k=20)
    outs = get_topk_from_heatmap__default(scores, k=20)

    for gt, out in zip(gts, outs):
        torch.testing.assert_close(gt, out)
//  [python]
def testRNNCellForward(self, _, cell_constructor, output_selector):
        num_layers = 2
        hidden_size = 20
        num_feats = 10

        batch_size = 3
        seq_length = 5

        for _ in range(10):
            cell = cell_constructor(num_feats, hidden_size, num_layers)
            fcell = higher.patch.monkeypatch(cell)

            state = fstate = None
            for _ in range(seq_length):
                inputs = torch.randn(batch_size, num_feats)

                state = cell(inputs, state)
                fstate = fcell(inputs, fstate)

                if isinstance(state, tuple):
                    self.assertEqual(len(state), len(fstate))
                    for s, fs in zip(state, fstate):
                        torch.testing.assert_allclose(s, fs)
                else:
                    torch.testing.assert_allclose(state, fstate)
//  [python]
def testRNNCellForward(self, _, cell_constructor, output_selector):
        num_layers = 2
        hidden_size = 20
        num_feats = 10

        batch_size = 3
        seq_length = 5

        for _ in range(10):
            cell = cell_constructor(num_feats, hidden_size, num_layers)
            fcell = higher.patch.monkeypatch(cell)

            state = fstate = None
            for _ in range(seq_length):
                inputs = torch.randn(batch_size, num_feats)

                state = cell(inputs, state)
                fstate = fcell(inputs, fstate)

                if isinstance(state, tuple):
                    self.assertEqual(len(state), len(fstate))
                    for s, fs in zip(state, fstate):
                        torch.testing.assert_close(s, fs)
                else:
                    torch.testing.assert_close(state, fstate)
//  [python]
def testMiniMAML(self, _, model_builder):
        model = model_builder(self)
        opt = torch.optim.SGD(model.parameters(), lr=self.lr)

        ctx_opts = {"copy_initial_weights": False}
        with higher.innerloop_ctx(model, opt, **ctx_opts) as (fmodel, diffopt):
            for _ in range(10):
                inputs = torch.rand(8, 4)
                loss = fmodel(inputs).sum().pow(2)
                diffopt.step(loss)
            param_sum = sum(p.sum() for p in fmodel.parameters())
            final_grads = torch.autograd.grad(
                param_sum, fmodel.parameters(time=0), retain_graph=True
            )
            param_sum.backward()
            for p, g in zip(model.parameters(), final_grads):
                torch.testing.assert_allclose(p.grad, g)
//  [python]
def testMiniMAML(self, _, model_builder):
        model = model_builder(self)
        opt = torch.optim.SGD(model.parameters(), lr=self.lr)

        ctx_opts = {"copy_initial_weights": False}
        with higher.innerloop_ctx(model, opt, **ctx_opts) as (fmodel, diffopt):
            for _ in range(10):
                inputs = torch.rand(8, 4)
                loss = fmodel(inputs).sum().pow(2)
                diffopt.step(loss)
            param_sum = sum(p.sum() for p in fmodel.parameters())
            final_grads = torch.autograd.grad(
                param_sum, fmodel.parameters(time=0), retain_graph=True
            )
            param_sum.backward()
            for p, g in zip(model.parameters(), final_grads):
                torch.testing.assert_close(p.grad, g)
//  [python]
def test_audio_samples_with_custom_field(libri_cut_set):
    batchio = AudioSamples()

    def attach_custom_audio(cut):
        """Simulate adding an additional custom recording"""
        cut.my_favorite_song = cut.recording.perturb_volume(factor=1.1)
        return cut

    # Reject mixed cuts (we don't support mixing custom attributes for now) and add custom audio
    cuts = libri_cut_set.filter(lambda c: isinstance(c, MonoCut)).map(
        attach_custom_audio
    )
    # does not fail / hang / etc.
    audio, audio_lens = batchio(cuts, recording_field="my_favorite_song")
    assert audio.shape[0] == 3

    # check that the audio is not the same as in the default 'recording' field
    audio_default, _ = batchio(cuts)
    with pytest.raises(AssertionError):
        torch.testing.assert_allclose(audio, audio_default)
//  [python]
def test_audio_samples_with_custom_field(libri_cut_set):
    batchio = AudioSamples()

    def attach_custom_audio(cut):
        """Simulate adding an additional custom recording"""
        cut.my_favorite_song = cut.recording.perturb_volume(factor=1.1)
        return cut

    # Reject mixed cuts (we don't support mixing custom attributes for now) and add custom audio
    cuts = libri_cut_set.filter(lambda c: isinstance(c, MonoCut)).map(
        attach_custom_audio
    )
    # does not fail / hang / etc.
    audio, audio_lens = batchio(cuts, recording_field="my_favorite_song")
    assert audio.shape[0] == 3

    # check that the audio is not the same as in the default 'recording' field
    audio_default, _ = batchio(cuts)
    with pytest.raises(AssertionError):
        torch.testing.assert_close(audio, audio_default)
//  [python]
def test_octconv_block(use_multi):
    block = OctConvBlock(3, 6, (0.0, 0.5), use_multi=use_multi)
    block.eval()  # Switch the model to inference mode

    # N, C, H, W
    x = torch.randn(5, 3, 10, 10)
    out = block(x)

    pop_block = poptorch.inferenceModel(block)
    pop_out = pop_block(x)

    for (cpu, ipu) in zip(out, pop_out):
        torch.testing.assert_allclose(cpu, ipu)
//  [python]
def test_octconv_block(use_multi):
    block = OctConvBlock(3, 6, (0.0, 0.5), use_multi=use_multi)
    block.eval()  # Switch the model to inference mode

    # N, C, H, W
    x = torch.randn(5, 3, 10, 10)
    out = block(x)

    pop_block = poptorch.inferenceModel(block)
    pop_out = pop_block(x)

    for (cpu, ipu) in zip(out, pop_out):
        torch.testing.assert_close(cpu, ipu)
//  [python]
def test_inference_model(conv_mode):
    model = ClassificationModel(conv_mode)
    model.eval()  # Switch the model to inference mode

    # N, C, H, W
    x = torch.randn(5, 3, 32, 32)
    out = model(x)

    pop_model = poptorch.inferenceModel(model)
    pop_out = pop_model(x)
    torch.testing.assert_allclose(out, pop_out)
//  [python]
def test_inference_model(conv_mode):
    model = ClassificationModel(conv_mode)
    model.eval()  # Switch the model to inference mode

    # N, C, H, W
    x = torch.randn(5, 3, 32, 32)
    out = model(x)

    pop_model = poptorch.inferenceModel(model)
    pop_out = pop_model(x)
    torch.testing.assert_close(out, pop_out)
//  [python]
def test_training_model(conv_mode):
    model = ClassificationModel(conv_mode)

    # N, C, H, W
    x = torch.randn(5, 3, 32, 32)
    labels = torch.randint(low=1, high=10, size=(5,))
    out, loss = model(x, labels)

    model.train()  # Switch the model to training mode
    # Models are initialised in training mode by default, so the line above will
    # have no effect. Its purpose is to show how the mode can be set explicitly.

    pop_model = poptorch.trainingModel(model, poptorch.Options(), poptorch.optim.SGD(model.parameters(), lr=0.01))
    pop_out, pop_loss = pop_model(x, labels)
    torch.testing.assert_allclose(out, pop_out)
    torch.testing.assert_allclose(loss, pop_loss)
//  [python]
def test_training_model(conv_mode):
    model = ClassificationModel(conv_mode)

    # N, C, H, W
    x = torch.randn(5, 3, 32, 32)
    labels = torch.randint(low=1, high=10, size=(5,))
    out, loss = model(x, labels)

    model.train()  # Switch the model to training mode
    # Models are initialised in training mode by default, so the line above will
    # have no effect. Its purpose is to show how the mode can be set explicitly.

    pop_model = poptorch.trainingModel(model, poptorch.Options(), poptorch.optim.SGD(model.parameters(), lr=0.01))
    pop_out, pop_loss = pop_model(x, labels)
    torch.testing.assert_close(out, pop_out)
    torch.testing.assert_close(loss, pop_loss)
//  [python]
def test_compute_loss():
    model = DummyModel((3, 3), 10)

    with pytest.raises(
        EngineException,
        match=r'Exception "ValueError\(\'The loss is NaN or ± inf\'[,]?\)" raised',
    ):
        module = EngineModule(model, lambda *_: torch.tensor([1, float("nan")]))
        module.compute_loss(None, None, None)

    expected = torch.tensor(0.1)
    module = EngineModule(model, lambda *_: expected)
    actual = module.compute_loss(None, None, None)
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_compute_loss():
    model = DummyModel((3, 3), 10)

    with pytest.raises(
        EngineException,
        match=r'Exception "ValueError\(\'The loss is NaN or ± inf\'[,]?\)" raised',
    ):
        module = EngineModule(model, lambda *_: torch.tensor([1, float("nan")]))
        module.compute_loss(None, None, None)

    expected = torch.tensor(0.1)
    module = EngineModule(model, lambda *_: expected)
    actual = module.compute_loss(None, None, None)
    torch.testing.assert_close(actual, expected)
//  [python]
def ort_validate(self, onnx_io, inputs, outputs):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        y_runtime = PredictorORT(onnx_io.getvalue())
        # Inference on ONNX Runtime
        ort_outs = y_runtime.predict(inputs)

        for i in range(0, len(outputs)):
            torch.testing.assert_allclose(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
//  [python]
def ort_validate(self, onnx_io, inputs, outputs):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        y_runtime = PredictorORT(onnx_io.getvalue())
        # Inference on ONNX Runtime
        ort_outs = y_runtime.predict(inputs)

        for i in range(0, len(outputs)):
            torch.testing.assert_close(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
//  [python]
def test_arrow_dataset_reload():
    """
    Validates that the dataset can be reloaded from a saved state dict.
    """
    with tempfile.TemporaryDirectory() as data_path:
        generate_test_data(data_path, 8, 97)
        data_path = "file:///" + data_path

        ds = ArrowFilesDataset(data_path)
        it = iter(ds)
        for i, batch in enumerate(it):
            assert batch[0][0] == i
            if i > 100:
                break

        sd = ds.state_dict()
        # print(sd)
        # save some data from after where it was saved, so we can verify
        # that it gets the same data after loading the sd
        saved = []
        for i, batch in enumerate(it):
            # print(i, batch)
            saved.append(batch[0])
            if i > 5:
                break

        # print("saved", saved)
        ds = ArrowFilesDataset(data_path)
        ds.load_state_dict(sd)
        # print("loaded state dict", ds.state_dict())
        for i, batch in enumerate(ds):
            torch.testing.assert_allclose(saved[i], batch[0])
            if i > 5:
                break
//  [python]
def test_arrow_dataset_reload():
    """
    Validates that the dataset can be reloaded from a saved state dict.
    """
    with tempfile.TemporaryDirectory() as data_path:
        generate_test_data(data_path, 8, 97)
        data_path = "file:///" + data_path

        ds = ArrowFilesDataset(data_path)
        it = iter(ds)
        for i, batch in enumerate(it):
            assert batch[0][0] == i
            if i > 100:
                break

        sd = ds.state_dict()
        # print(sd)
        # save some data from after where it was saved, so we can verify
        # that it gets the same data after loading the sd
        saved = []
        for i, batch in enumerate(it):
            # print(i, batch)
            saved.append(batch[0])
            if i > 5:
                break

        # print("saved", saved)
        ds = ArrowFilesDataset(data_path)
        ds.load_state_dict(sd)
        # print("loaded state dict", ds.state_dict())
        for i, batch in enumerate(ds):
            torch.testing.assert_close(saved[i], batch[0])
            if i > 5:
                break
//  [python]
def test_transform_batch():
    with pytest.raises(
        NotImplementedError, match=r"Not implemented for type <class 'NoneType'>"
    ):
        transform_batch(None)

    batch = torch.empty((3, 4, 5))
    x, xs = transform_batch(batch)
    torch.testing.assert_allclose(x, batch)
    assert xs == [3, 3, 3, 3]

    x = torch.tensor([[4, 5, 6], [1, 2, 0], [3, 0, 0]])
    xs = [3, 2, 1]
    batch = torch.nn.utils.rnn.pack_padded_sequence(x, xs)
    x_out, xs_out = transform_batch(batch)
    torch.testing.assert_allclose(x_out, x)
    assert xs_out.tolist() == xs
//  [python]
def test_transform_batch():
    with pytest.raises(
        NotImplementedError, match=r"Not implemented for type <class 'NoneType'>"
    ):
        transform_batch(None)

    batch = torch.empty((3, 4, 5))
    x, xs = transform_batch(batch)
    torch.testing.assert_close(x, batch)
    assert xs == [3, 3, 3, 3]

    x = torch.tensor([[4, 5, 6], [1, 2, 0], [3, 0, 0]])
    xs = [3, 2, 1]
    batch = torch.nn.utils.rnn.pack_padded_sequence(x, xs)
    x_out, xs_out = transform_batch(batch)
    torch.testing.assert_close(x_out, x)
    assert xs_out.tolist() == xs
//  [python]
def test_ranked():
    """
    test handling of rank and world size in saved state dict.
    """
    with tempfile.TemporaryDirectory() as data_path:
        generate_test_data(data_path, 7, 95)
        data_path = "file:///" + data_path

        rank = 7
        world = 13

        # we actually save the sd from rank 0. other ranks take every
        # rank'th row after the zeroth row.
        ds0 = ArrowFilesDataset(data_path, 0, world)
        ds = ArrowFilesDataset(data_path, rank, world)

        it = iter(ds)
        iter0 = iter(ds0)
        for i, batch in enumerate(it):
            next(iter0)
            v = batch[0][0]
            assert v % world == rank
            if v > 250:
                break

        sd = ds0.state_dict()
        print(sd)
        saved = []
        for i, batch in enumerate(it):
            saved.append(batch[0])
            if i > 5:
                break

        ds = ArrowFilesDataset(data_path, rank, world)
        ds.load_state_dict(sd)
        for i, batch in enumerate(ds):
            # print(i, batch, saved[i])
            torch.testing.assert_allclose(saved[i], batch[0])
            if i > 5:
                break
//  [python]
def test_ranked():
    """
    test handling of rank and world size in saved state dict.
    """
    with tempfile.TemporaryDirectory() as data_path:
        generate_test_data(data_path, 7, 95)
        data_path = "file:///" + data_path

        rank = 7
        world = 13

        # we actually save the sd from rank 0. other ranks take every
        # rank'th row after the zeroth row.
        ds0 = ArrowFilesDataset(data_path, 0, world)
        ds = ArrowFilesDataset(data_path, rank, world)

        it = iter(ds)
        iter0 = iter(ds0)
        for i, batch in enumerate(it):
            next(iter0)
            v = batch[0][0]
            assert v % world == rank
            if v > 250:
                break

        sd = ds0.state_dict()
        print(sd)
        saved = []
        for i, batch in enumerate(it):
            saved.append(batch[0])
            if i > 5:
                break

        ds = ArrowFilesDataset(data_path, rank, world)
        ds.load_state_dict(sd)
        for i, batch in enumerate(ds):
            # print(i, batch, saved[i])
            torch.testing.assert_close(saved[i], batch[0])
            if i > 5:
                break
//  [python]
def test_split():
    """
    in cases where a row is longer than max seq len, it's chunked and
    read as chunks. this test checks the handling of the state dict for
    picking up after the last chunk.
    """
    with tempfile.TemporaryDirectory() as data_path:
        generate_test_data(data_path, 7, 10)
        data_path = "file:///" + data_path

        max_seq_len = 13
        ds = ArrowFilesDataset(data_path, max_seq_len=max_seq_len)

        it = iter(ds)
        for i, batch in enumerate(it):
            if i > 100:
                break

        sd = ds.state_dict()
        print(sd)
        saved = []
        for i, batch in enumerate(it):
            # print(i, batch)
            saved.append(batch)
            if i > 5:
                break

        # print("saved", saved)
        ds = ArrowFilesDataset(data_path, max_seq_len=max_seq_len)
        ds.load_state_dict(sd)
        # print("loaded state dict", ds.state_dict())
        for i, batch in enumerate(ds):
            torch.testing.assert_allclose(saved[i][0], batch[0])
            torch.testing.assert_allclose(saved[i][1], batch[1])
            if i > 5:
                break
//  [python]
def test_split():
    """
    in cases where a row is longer than max seq len, it's chunked and
    read as chunks. this test checks the handling of the state dict for
    picking up after the last chunk.
    """
    with tempfile.TemporaryDirectory() as data_path:
        generate_test_data(data_path, 7, 10)
        data_path = "file:///" + data_path

        max_seq_len = 13
        ds = ArrowFilesDataset(data_path, max_seq_len=max_seq_len)

        it = iter(ds)
        for i, batch in enumerate(it):
            if i > 100:
                break

        sd = ds.state_dict()
        print(sd)
        saved = []
        for i, batch in enumerate(it):
            # print(i, batch)
            saved.append(batch)
            if i > 5:
                break

        # print("saved", saved)
        ds = ArrowFilesDataset(data_path, max_seq_len=max_seq_len)
        ds.load_state_dict(sd)
        # print("loaded state dict", ds.state_dict())
        for i, batch in enumerate(ds):
            torch.testing.assert_close(saved[i][0], batch[0])
            torch.testing.assert_close(saved[i][1], batch[1])
            if i > 5:
                break
//  [python]
def test_compute(self):
        from latte.metrics.torch.wrapper import TorchMetricWrapper

        val = np.random.randn(16,)
        dummy_metric = TorchMetricWrapper(metric=DummyMetric, val=val)

        newval = torch.randn(size=(16,))
        mult = torch.randn(size=(1,))
        dummy_metric.update(newval, mult=mult)

        out = dummy_metric.compute()

        assert isinstance(out, torch.Tensor)
        torch.testing.assert_allclose(out, 2.0 * newval * mult)
//  [python]
def test_compute(self):
        from latte.metrics.torch.wrapper import TorchMetricWrapper

        val = np.random.randn(16,)
        dummy_metric = TorchMetricWrapper(metric=DummyMetric, val=val)

        newval = torch.randn(size=(16,))
        mult = torch.randn(size=(1,))
        dummy_metric.update(newval, mult=mult)

        out = dummy_metric.compute()

        assert isinstance(out, torch.Tensor)
        torch.testing.assert_close(out, 2.0 * newval * mult)
//  [python]
def main(cfg):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    OmegaConf.set_struct(cfg, False)
    cfg.pop("ckpt")
    OmegaConf.set_struct(cfg, True)

    model = MineCLIP(**cfg).to(device)

    video = torch.randint(0, 255, (6, 16, 3, 160, 256), device=device)
    prompts = [
        "hello, this is MineCLIP",
        "MineCLIP is a VideoCLIP model trained on YouTube dataset from MineDojo's knowledge base",
        "Feel free to also checkout MineDojo at",
        "https://minedojo.org",
    ]
    VIDEO_BATCH, TEXT_BATCH = video.size(0), len(prompts)

    image_feats = model.forward_image_features(video)
    video_feats = model.forward_video_features(image_feats)
    assert video_feats.shape == (VIDEO_BATCH, 512)
    video_feats_2 = model.encode_video(video)
    # encode_video is equivalent to forward_video_features(forward_image_features(video))
    torch.testing.assert_allclose(video_feats, video_feats_2)

    # encode batch of prompts
    text_feats_batch = model.encode_text(prompts)
    assert text_feats_batch.shape == (TEXT_BATCH, 512)

    # compute reward from features
    logits_per_video, logits_per_text = model.forward_reward_head(
        video_feats, text_tokens=text_feats_batch
    )
    assert logits_per_video.shape == (VIDEO_BATCH, TEXT_BATCH)
    assert logits_per_text.shape == (TEXT_BATCH, VIDEO_BATCH)
    # directly pass in strings. This invokes the tokenizer under the hood
    reward_scores_2, _ = model.forward_reward_head(video_feats, text_tokens=prompts)
    # pass in cached, encoded text features
    reward_scores_3, _ = model(
        video_feats, text_tokens=text_feats_batch, is_video_features=True
    )
    reward_scores_4, _ = model(
        video, text_tokens=text_feats_batch, is_video_features=False
    )
    # all above are equivalent, just starting from features or raw values
    torch.testing.assert_allclose(logits_per_video, reward_scores_2)
    torch.testing.assert_allclose(logits_per_video, reward_scores_3)
    torch.testing.assert_allclose(logits_per_video, reward_scores_4)

    print("Inference successful")
//  [python]
def main(cfg):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    OmegaConf.set_struct(cfg, False)
    cfg.pop("ckpt")
    OmegaConf.set_struct(cfg, True)

    model = MineCLIP(**cfg).to(device)

    video = torch.randint(0, 255, (6, 16, 3, 160, 256), device=device)
    prompts = [
        "hello, this is MineCLIP",
        "MineCLIP is a VideoCLIP model trained on YouTube dataset from MineDojo's knowledge base",
        "Feel free to also checkout MineDojo at",
        "https://minedojo.org",
    ]
    VIDEO_BATCH, TEXT_BATCH = video.size(0), len(prompts)

    image_feats = model.forward_image_features(video)
    video_feats = model.forward_video_features(image_feats)
    assert video_feats.shape == (VIDEO_BATCH, 512)
    video_feats_2 = model.encode_video(video)
    # encode_video is equivalent to forward_video_features(forward_image_features(video))
    torch.testing.assert_close(video_feats, video_feats_2)

    # encode batch of prompts
    text_feats_batch = model.encode_text(prompts)
    assert text_feats_batch.shape == (TEXT_BATCH, 512)

    # compute reward from features
    logits_per_video, logits_per_text = model.forward_reward_head(
        video_feats, text_tokens=text_feats_batch
    )
    assert logits_per_video.shape == (VIDEO_BATCH, TEXT_BATCH)
    assert logits_per_text.shape == (TEXT_BATCH, VIDEO_BATCH)
    # directly pass in strings. This invokes the tokenizer under the hood
    reward_scores_2, _ = model.forward_reward_head(video_feats, text_tokens=prompts)
    # pass in cached, encoded text features
    reward_scores_3, _ = model(
        video_feats, text_tokens=text_feats_batch, is_video_features=True
    )
    reward_scores_4, _ = model(
        video, text_tokens=text_feats_batch, is_video_features=False
    )
    # all above are equivalent, just starting from features or raw values
    torch.testing.assert_close(logits_per_video, reward_scores_2)
    torch.testing.assert_close(logits_per_video, reward_scores_3)
    torch.testing.assert_close(logits_per_video, reward_scores_4)

    print("Inference successful")
//  [python]
def test_load_save_csuite(dataset):
    """Test that we can load a csuite dataset, save it, and load it again"""
    root_path = os.path.join(CAUSICA_DATASETS_PATH, dataset)

    variables_metadata, train_data, test_data, adj_mat, interventions = _load_and_test_dataset(root_path)

    with tempfile.TemporaryDirectory() as tmpdir:
        save_dataset(
            tmpdir, variables_metadata, adj_mat, train_data, test_data, interventions=interventions, overwrite=True
        )
        variables_metadata_2, train_data_2, test_data_2, adj_mat_2, interventions_2 = _load_and_test_dataset(tmpdir)

    assert variables_metadata == variables_metadata_2
    torch.testing.assert_allclose(adj_mat, adj_mat_2)
    torch.testing.assert_allclose(train_data, train_data_2)
    assert tensordict_shapes(train_data) == tensordict_shapes(train_data_2)
    assert tensordict_shapes(test_data) == tensordict_shapes(test_data_2)
    assert len(interventions) == len(interventions_2)
    for intervention_1, intervention_2 in zip(interventions, interventions_2):
        for field in ["intervention_data", "intervention_values", "condition_values"]:
            torch.testing.assert_allclose(getattr(intervention_1[0], field), getattr(intervention_2[0], field))
            torch.testing.assert_allclose(getattr(intervention_1[1], field), getattr(intervention_2[1], field))
        assert intervention_1[2] == intervention_2[2]
//  [python]
def test_load_save_csuite(dataset):
    """Test that we can load a csuite dataset, save it, and load it again"""
    root_path = os.path.join(CAUSICA_DATASETS_PATH, dataset)

    variables_metadata, train_data, test_data, adj_mat, interventions = _load_and_test_dataset(root_path)

    with tempfile.TemporaryDirectory() as tmpdir:
        save_dataset(
            tmpdir, variables_metadata, adj_mat, train_data, test_data, interventions=interventions, overwrite=True
        )
        variables_metadata_2, train_data_2, test_data_2, adj_mat_2, interventions_2 = _load_and_test_dataset(tmpdir)

    assert variables_metadata == variables_metadata_2
    torch.testing.assert_close(adj_mat, adj_mat_2)
    torch.testing.assert_close(train_data, train_data_2)
    assert tensordict_shapes(train_data) == tensordict_shapes(train_data_2)
    assert tensordict_shapes(test_data) == tensordict_shapes(test_data_2)
    assert len(interventions) == len(interventions_2)
    for intervention_1, intervention_2 in zip(interventions, interventions_2):
        for field in ["intervention_data", "intervention_values", "condition_values"]:
            torch.testing.assert_close(getattr(intervention_1[0], field), getattr(intervention_2[0], field))
            torch.testing.assert_close(getattr(intervention_1[1], field), getattr(intervention_2[1], field))
        assert intervention_1[2] == intervention_2[2]
//  [python]
def test_load_save_counterfactuals():
    dataset = "csuite_linexp_2"
    root_path = os.path.join(CAUSICA_DATASETS_PATH, dataset)

    variables_metadata, counterfactuals = load_and_test_counterfactuals(root_path)

    with tempfile.TemporaryDirectory() as tmpdir:
        save_data(tmpdir, variables_metadata, DataEnum.VARIABLES_JSON, variables_metadata)
        save_data(tmpdir, counterfactuals, DataEnum.COUNTERFACTUALS, variables_metadata)
        variables_metadata_2, counterfactuals_2 = load_and_test_counterfactuals(tmpdir)

    assert variables_metadata == variables_metadata_2
    assert len(counterfactuals) == len(counterfactuals_2)
    for cf_1, cf_2 in zip(counterfactuals, counterfactuals_2):
        for field in ["counterfactual_data", "intervention_values", "factual_data"]:
            torch.testing.assert_allclose(getattr(cf_1[0], field), getattr(cf_2[0], field))
            torch.testing.assert_allclose(getattr(cf_1[1], field), getattr(cf_2[1], field))
        assert cf_1[2] == cf_2[2]
//  [python]
def test_load_save_counterfactuals():
    dataset = "csuite_linexp_2"
    root_path = os.path.join(CAUSICA_DATASETS_PATH, dataset)

    variables_metadata, counterfactuals = load_and_test_counterfactuals(root_path)

    with tempfile.TemporaryDirectory() as tmpdir:
        save_data(tmpdir, variables_metadata, DataEnum.VARIABLES_JSON, variables_metadata)
        save_data(tmpdir, counterfactuals, DataEnum.COUNTERFACTUALS, variables_metadata)
        variables_metadata_2, counterfactuals_2 = load_and_test_counterfactuals(tmpdir)

    assert variables_metadata == variables_metadata_2
    assert len(counterfactuals) == len(counterfactuals_2)
    for cf_1, cf_2 in zip(counterfactuals, counterfactuals_2):
        for field in ["counterfactual_data", "intervention_values", "factual_data"]:
            torch.testing.assert_close(getattr(cf_1[0], field), getattr(cf_2[0], field))
            torch.testing.assert_close(getattr(cf_1[1], field), getattr(cf_2[1], field))
        assert cf_1[2] == cf_2[2]
//  [python]
def test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value(self):
        group, group_id, rank = self._init_global_test()
        rank_to_GPU = self._init_multigpu_helper()
        model = nn.parallel.DistributedDataParallel(ONLY_SBN_NET.cuda(rank), device_ids=[rank])

        input_var = []
        for i in range(int(WORLD_SIZE)):
            input_var_rank = torch.cat([
                torch.ones(2, 1, 10 ** (i + 1)) * (0.1 ** (i - 1)),
                torch.ones(2, 1, 10 ** (i + 1)) * (0.3 ** (i - 1))
            ], dim=1)
            input_var.append(input_var_rank)

        all_input_var = torch.cat(
            [x.permute(1, 0, 2).contiguous().view(ONLY_SBN_NET.num_features, -1) for x in input_var],
            dim=1
        ).cuda(rank)

        for i in range(100):
            y = model(input_var[rank].cuda(rank))
            y.mean().backward()

        running_mean, running_var = model.module.running_mean, model.module.running_var
        torch.testing.assert_allclose(running_mean, all_input_var.mean(1))
        torch.testing.assert_allclose(running_var, all_input_var.var(1))
//  [python]
def test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value(self):
        group, group_id, rank = self._init_global_test()
        rank_to_GPU = self._init_multigpu_helper()
        model = nn.parallel.DistributedDataParallel(ONLY_SBN_NET.cuda(rank), device_ids=[rank])

        input_var = []
        for i in range(int(WORLD_SIZE)):
            input_var_rank = torch.cat([
                torch.ones(2, 1, 10 ** (i + 1)) * (0.1 ** (i - 1)),
                torch.ones(2, 1, 10 ** (i + 1)) * (0.3 ** (i - 1))
            ], dim=1)
            input_var.append(input_var_rank)

        all_input_var = torch.cat(
            [x.permute(1, 0, 2).contiguous().view(ONLY_SBN_NET.num_features, -1) for x in input_var],
            dim=1
        ).cuda(rank)

        for i in range(100):
            y = model(input_var[rank].cuda(rank))
            y.mean().backward()

        running_mean, running_var = model.module.running_mean, model.module.running_var
        torch.testing.assert_close(running_mean, all_input_var.mean(1))
        torch.testing.assert_close(running_var, all_input_var.var(1))
//  [python]
def test_no_register_hook(self):
        model_with_hook = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 10))
        initial_model = deepcopy(model_with_hook)
        model_no_hook = deepcopy(model_with_hook)
        _apply_optimizer_in_backward(
            torch.optim.SGD,
            model_with_hook.parameters(),
            optimizer_kwargs={"lr": 0.03},
        )
        _apply_optimizer_in_backward(
            torch.optim.SGD,
            model_no_hook.parameters(),
            optimizer_kwargs={"lr": 0.03},
            register_hook=False,
        )
        inp = torch.randn(4, 10)
        model_with_hook(inp).sum().backward()
        model_no_hook(inp).sum().backward()

        for p1, p2 in zip(model_with_hook.parameters(), initial_model.parameters()):
            with self.assertRaises(AssertionError):
                torch.testing.assert_allclose(p1, p2)

        for p1, p2 in zip(model_no_hook.parameters(), initial_model.parameters()):
            torch.testing.assert_allclose(p1, p2)
//  [python]
def test_no_register_hook(self):
        model_with_hook = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 10))
        initial_model = deepcopy(model_with_hook)
        model_no_hook = deepcopy(model_with_hook)
        _apply_optimizer_in_backward(
            torch.optim.SGD,
            model_with_hook.parameters(),
            optimizer_kwargs={"lr": 0.03},
        )
        _apply_optimizer_in_backward(
            torch.optim.SGD,
            model_no_hook.parameters(),
            optimizer_kwargs={"lr": 0.03},
            register_hook=False,
        )
        inp = torch.randn(4, 10)
        model_with_hook(inp).sum().backward()
        model_no_hook(inp).sum().backward()

        for p1, p2 in zip(model_with_hook.parameters(), initial_model.parameters()):
            with self.assertRaises(AssertionError):
                torch.testing.assert_close(p1, p2)

        for p1, p2 in zip(model_no_hook.parameters(), initial_model.parameters()):
            torch.testing.assert_close(p1, p2)