jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_initialization(self) -> None:
        in_features = 10
        out_features = 20
        biases = [False, True]
        seeds = [920, 2022]

        for bias, seed in itertools.product(biases, seeds):
            with self.subTest(
                f"Test with in_features = {in_features}, out_features = {out_features}, bias = {bias}, seed = {seed}"
            ):
                # generate linear layers and initialize
                torch.random.manual_seed(seed)
                original_linear = nn.Linear(in_features, out_features, bias=bias)
                torch.random.manual_seed(seed)
                combined_linear = CombinedLinear(in_features, out_features, bias=bias)

                # confirm weights are initialized equally
                if bias:
                    self.assertEqual(original_linear.weight.shape, combined_linear.combined_weight[:, :-1].shape)
                    self.assertTrue(torch.allclose(original_linear.weight, combined_linear.combined_weight[:, :-1]))
                    self.assertEqual(original_linear.bias.shape, combined_linear.combined_weight[:, -1].shape)
                    self.assertTrue(torch.allclose(original_linear.bias, combined_linear.combined_weight[:, -1]))
                else:
                    self.assertEqual(original_linear.weight.shape, combined_linear.combined_weight.shape)
                    self.assertTrue(torch.allclose(original_linear.weight, combined_linear.combined_weight))
//  [python]
def test_initialization(self) -> None:
        in_features = 10
        out_features = 20
        biases = [False, True]
        seeds = [920, 2022]

        for bias, seed in itertools.product(biases, seeds):
            with self.subTest(
                f"Test with in_features = {in_features}, out_features = {out_features}, bias = {bias}, seed = {seed}"
            ):
                # generate linear layers and initialize
                torch.random.manual_seed(seed)
                original_linear = nn.Linear(in_features, out_features, bias=bias)
                torch.random.manual_seed(seed)
                combined_linear = CombinedLinear(in_features, out_features, bias=bias)

                # confirm weights are initialized equally
                if bias:
                    torch.testing.assert_close(
                        original_linear.weight, combined_linear.combined_weight[:, :-1]
                    )
                    torch.testing.assert_close(
                        original_linear.bias, combined_linear.combined_weight[:, -1]
                    )
                else:
                    torch.testing.assert_close(
                        original_linear.weight, combined_linear.combined_weight
                    )
//  [python]
def _test_linear_forward_backward(
        self,
        feature_vector: torch.Tensor,
        in_features: int,
        out_features: int,
        bias: bool,
        seed: int,
    ) -> None:
        # generate linear layers and initialize
        original_linear = nn.Linear(in_features, out_features, bias=bias)
        combined_linear = CombinedLinear(in_features, out_features, bias=bias)

        self._init_weights(original_linear, seed)
        self._init_weights(combined_linear, seed)

        # confirm weights are initialized equally
        if bias:
            assert torch.equal(
                original_linear.weight, combined_linear.combined_weight[:, :-1]
            )
            assert torch.equal(
                original_linear.bias, combined_linear.combined_weight[:, -1]
            )
        else:
            assert torch.equal(original_linear.weight, combined_linear.combined_weight)

        # perform forward pass
        original_output = original_linear(feature_vector)
        combined_output = combined_linear(feature_vector)

        # compute backward of sum of output
        torch.sum(original_output).backward()
        torch.sum(combined_output).backward()

        # check values are equal
        with self.subTest("Test forward"):
            t.assert_close(original_output, combined_output)

        with self.subTest("Test backward"):
            if bias:
                t.assert_close(
                    cast(torch.Tensor, original_linear.weight.grad),
                    cast(torch.Tensor, combined_linear.combined_weight.grad)[:, :-1],
                )
                t.assert_close(
                    cast(torch.Tensor, original_linear.bias.grad),
                    cast(torch.Tensor, combined_linear.combined_weight.grad)[:, -1],
                )
            else:
                t.assert_close(
                    cast(torch.Tensor, original_linear.weight.grad),
                    cast(torch.Tensor, combined_linear.combined_weight.grad),
                )
//  [python]
def _test_linear_forward_backward(
        self,
        feature_vector: torch.Tensor,
        in_features: int,
        out_features: int,
        bias: bool,
        seed: int,
    ) -> None:
        # generate linear layers and initialize
        original_linear = nn.Linear(in_features, out_features, bias=bias)
        combined_linear = CombinedLinear(in_features, out_features, bias=bias)

        self._init_weights(original_linear, seed)
        self._init_weights(combined_linear, seed)

        # confirm weights are initialized equally
        if bias:
            assert torch.equal(
                original_linear.weight, combined_linear.combined_weight[:, :-1]
            )
            assert torch.equal(
                original_linear.bias, combined_linear.combined_weight[:, -1]
            )
        else:
            assert torch.equal(original_linear.weight, combined_linear.combined_weight)

        # perform forward pass
        original_output = original_linear(feature_vector)
        combined_output = combined_linear(feature_vector)

        # compute backward of sum of output
        torch.sum(original_output).backward()
        torch.sum(combined_output).backward()

        # check values are equal
        with self.subTest("Test forward"):
            torch.testing.assert_close(original_output, combined_output)

        with self.subTest("Test backward"):
            if bias:
                torch.testing.assert_close(
                    cast(torch.Tensor, original_linear.weight.grad),
                    cast(torch.Tensor, combined_linear.combined_weight.grad)[:, :-1],
                )
                torch.testing.assert_close(
                    cast(torch.Tensor, original_linear.bias.grad),
                    cast(torch.Tensor, combined_linear.combined_weight.grad)[:, -1],
                )
            else:
                torch.testing.assert_close(
                    cast(torch.Tensor, original_linear.weight.grad),
                    cast(torch.Tensor, combined_linear.combined_weight.grad),
                )
//  [python]
def test_cached_layer_dict(self):
        """Check cached layer saves dict"""
        model = DivideInputBy2OutputDict()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = [torch.randn(1) for _ in range(2)]
        output = model(input)
        torch.testing.assert_allclose(output, cache["test_layer"])
//  [python]
def test_cached_layer_dict(self):
        """Check cached layer saves dict"""
        model = DivideInputBy2OutputDict()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = [torch.randn(1) for _ in range(2)]
        output = model(input)
        torch.testing.assert_close(output, cache["test_layer"])
//  [python]
def test_cached_layer_arbitrary(self):
        """Check cached layer saves arbitrary nested data structure"""
        model = TimesTable5OutputDict()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = [torch.randn(1) for _ in range(2)]
        output = model(input)
        t.assert_close(output, cache["test_layer"])
//  [python]
def test_cached_layer_arbitrary(self):
        """Check cached layer saves arbitrary nested data structure"""
        model = TimesTable5OutputDict()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = [torch.randn(1) for _ in range(2)]
        output = model(input)
        torch.testing.assert_close(output, cache["test_layer"])
//  [python]
def test_noop(self):
        """Check noop"""
        pseudo_labeler = NoopPseudoLabeler()
        x = np.random.randn(1)
        output = pseudo_labeler.label(x)
        t.assert_close(x, output)
//  [python]
def test_noop(self):
        """Check noop"""
        pseudo_labeler = NoopPseudoLabeler()
        x = np.random.randn(1)
        output = pseudo_labeler.label(x)
        torch.testing.assert_close(x, output)
//  [python]
def test_record_layers(self):
        """Check we can record specified layer"""
        model = AddLayers()
        cache = record_layers(model, ["", "layer0", "layer1", "layer2"])

        input = torch.Tensor([0])
        output = model(input)
        t.assert_close(cache["layer0"], torch.Tensor([1]))
        t.assert_close(cache["layer1"], torch.Tensor([2]))
        t.assert_close(cache["layer2"], torch.Tensor([3]))
        t.assert_close(cache[""], output)
//  [python]
def test_record_layers(self):
        """Check we can record specified layer"""
        model = AddLayers()
        cache = record_layers(model, ["", "layer0", "layer1", "layer2"])

        input = torch.Tensor([0])
        output = model(input)
        torch.testing.assert_close(cache["layer0"], torch.Tensor([1]))
        torch.testing.assert_close(cache["layer1"], torch.Tensor([2]))
        torch.testing.assert_close(cache["layer2"], torch.Tensor([3]))
        torch.testing.assert_close(cache[""], output)
//  [python]
def test_set_cache_dict(self):
        """Check we can swap the cache dict used when recording layers"""
        model = AddLayers()
        cache = record_layers(model, ["", "layer0", "layer1", "layer2"])
        new_cache = {}
        set_cache_dict(model, new_cache)
        input = torch.Tensor([0])
        output = model(input)
        self.assertEqual(cache, {})
        self.assertEqual(new_cache["layer0"], torch.Tensor([1]))
        self.assertEqual(new_cache["layer1"], torch.Tensor([2]))
        self.assertEqual(new_cache["layer2"], torch.Tensor([3]))
        self.assertEqual(new_cache[""], output)
//  [python]
def test_set_cache_dict(self):
        """Check we can swap the cache dict used when recording layers"""
        model = AddLayers()
        cache = record_layers(model, ["", "layer0", "layer1", "layer2"])
        new_cache = {}
        set_cache_dict(model, new_cache)
        input = torch.Tensor([0])
        output = model(input)
        self.assertEqual(cache, {})
        torch.testing.assert_close(new_cache["layer0"], torch.Tensor([1]))
        torch.testing.assert_close(new_cache["layer1"], torch.Tensor([2]))
        torch.testing.assert_close(new_cache["layer2"], torch.Tensor([3]))
        torch.testing.assert_close(new_cache[""], output)
//  [python]
def test_relabeltargetinbatch(self):
        """Check target is relabed using teacher"""
        teacher = DivideInputDictBy2()
        teacher.eval()
        teacher.device = torch.device("cpu")
        relabeler = RelabelTargetInBatch(teacher=teacher)
        batched_inputs = _get_input_data(n=2, use_input_target=True)
        gt = [{"input": d["input"], "target": d["input"] / 2.0} for d in batched_inputs]
        outputs = relabeler.label(batched_inputs)
        t.assert_close(outputs, gt)
//  [python]
def test_relabeltargetinbatch(self):
        """Check target is relabed using teacher"""
        teacher = DivideInputDictBy2()
        teacher.eval()
        teacher.device = torch.device("cpu")
        relabeler = RelabelTargetInBatch(teacher=teacher)
        batched_inputs = _get_input_data(n=2, use_input_target=True)
        gt = [{"input": d["input"], "target": d["input"] / 2.0} for d in batched_inputs]
        outputs = relabeler.label(batched_inputs)
        torch.testing.assert_close(outputs, gt)
//  [python]
def test_compute_layer_losses(self):
        """Check iterating over loss dicts"""
        layer_losses = [
            LayerLossMetadata(
                loss=lambda x, y: x + y, name="add", layer0="l00", layer1="l10"
            ),
            LayerLossMetadata(
                loss=lambda x, y: x / y, name="div", layer0="l01", layer1="l11"
            ),
        ]
        layer0_cache = {"l00": torch.randn(1), "l01": torch.randn(1)}
        layer1_cache = {"l10": torch.randn(1), "l11": torch.randn(1)}
        output = compute_layer_losses(layer_losses, layer0_cache, layer1_cache)
        t.assert_close(
            output["add"], layer0_cache["l00"] + layer1_cache["l10"]
        )
        t.assert_close(
            output["div"], layer0_cache["l01"] / layer1_cache["l11"]
        )
//  [python]
def test_compute_layer_losses(self):
        """Check iterating over loss dicts"""
        layer_losses = [
            LayerLossMetadata(
                loss=lambda x, y: x + y, name="add", layer0="l00", layer1="l10"
            ),
            LayerLossMetadata(
                loss=lambda x, y: x / y, name="div", layer0="l01", layer1="l11"
            ),
        ]
        layer0_cache = {"l00": torch.randn(1), "l01": torch.randn(1)}
        layer1_cache = {"l10": torch.randn(1), "l11": torch.randn(1)}
        output = compute_layer_losses(layer_losses, layer0_cache, layer1_cache)
        torch.testing.assert_close(
            output["add"], layer0_cache["l00"] + layer1_cache["l10"]
        )
        torch.testing.assert_close(
            output["div"], layer0_cache["l01"] / layer1_cache["l11"]
        )
//  [python]
def test_kd_inference(self):
        """Check inference defaults to student (and preprocessing)"""
        distillation_helper = TestHelper(cfg=CfgNode(), teacher=AddLayers())
        model = AddLayers()
        dynamic_mixin(
            model,
            KnowledgeDistillation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.eval()
        input = torch.randn(1)
        output = model(input)
        torch.testing.assert_allclose(output, input + 4.0)
//  [python]
def test_kd_inference(self):
        """Check inference defaults to student (and preprocessing)"""
        distillation_helper = TestHelper(cfg=CfgNode(), teacher=AddLayers())
        model = AddLayers()
        dynamic_mixin(
            model,
            KnowledgeDistillation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.eval()
        input = torch.randn(1)
        output = model(input)
        torch.testing.assert_close(output, input + 4.0)
//  [python]
def test_label_distillation_training(self):
        """Check training uses pseudo labeler

        Distillation teacher should run the teacher model on the inputs and
        then pass to the noop
        """
        batched_inputs = _get_input_data(n=2, requires_grad=True)
        gt = [x / 2.0 for x in batched_inputs]
        model = self.LabelDistillationNoop()
        model.dynamic_mixin_init(
            distillation_helper=TestHelper(cfg=None, teacher=DivideInputBy2()),
        )
        model.train()
        output = model(batched_inputs)
        torch.testing.assert_allclose(output, gt)

        sum(output).backward()
        torch.testing.assert_allclose(batched_inputs.grad, torch.Tensor([0.5, 0.5]))
//  [python]
def test_label_distillation_training(self):
        """Check training uses pseudo labeler

        Distillation teacher should run the teacher model on the inputs and
        then pass to the noop
        """
        batched_inputs = _get_input_data(n=2, requires_grad=True)
        gt = [x / 2.0 for x in batched_inputs]
        model = self.LabelDistillationNoop()
        model.dynamic_mixin_init(
            distillation_helper=TestHelper(cfg=None, teacher=DivideInputBy2()),
        )
        model.train()
        output = model(batched_inputs)
        torch.testing.assert_close(output, gt)

        sum(output).backward()
        torch.testing.assert_close(batched_inputs.grad, torch.Tensor([0.5, 0.5]))
//  [python]
def test_da_inference(self):
        """Check inference defaults to student (and preprocessing)"""
        distillation_helper = TestDAHelper(cfg=CfgNode(), teacher=nn.Identity())
        model = AddLayers()
        dynamic_mixin(
            model,
            DomainAdaptation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.eval()
        input = {"real": torch.randn(1), "synthetic": torch.randn(1)}
        output = model(input)
        t.assert_close(output, input["real"] + 3.0)
//  [python]
def test_da_inference(self):
        """Check inference defaults to student (and preprocessing)"""
        distillation_helper = TestDAHelper(cfg=CfgNode(), teacher=nn.Identity())
        model = AddLayers()
        dynamic_mixin(
            model,
            DomainAdaptation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.eval()
        input = {"real": torch.randn(1), "synthetic": torch.randn(1)}
        output = model(input)
        torch.testing.assert_close(output, input["real"] + 3.0)
//  [python]
def test_kd_train(self):
        """Check train pass results in updated loss output"""
        distillation_helper = TestHelper(cfg=CfgNode(), teacher=AddLayers())
        model = AddLayers()
        dynamic_mixin(
            model,
            KnowledgeDistillation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.train()
        input = torch.randn(1)
        output = model(input)
        torch.testing.assert_allclose(output["output"], (input + 4.0) * 0.1)
        torch.testing.assert_allclose(output["add"], ((input + 2.0) + (input + 3.0)) * 0.5)
        torch.testing.assert_allclose(output["mul"], (input + 3.0) * (input + 4.0) * 10.0)
//  [python]
def test_kd_train(self):
        """Check train pass results in updated loss output"""
        distillation_helper = TestHelper(cfg=CfgNode(), teacher=AddLayers())
        model = AddLayers()
        dynamic_mixin(
            model,
            KnowledgeDistillation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.train()
        input = torch.randn(1)
        output = model(input)
        torch.testing.assert_close(output["output"], (input + 4.0) * 0.1)
        torch.testing.assert_close(output["add"], ((input + 2.0) + (input + 3.0)) * 0.5)
        torch.testing.assert_close(output["mul"], (input + 3.0) * (input + 4.0) * 10.0)
//  [python]
def test_da_train(self):
        """Check train pass results in updated loss output"""
        distillation_helper = TestDAHelper(cfg=CfgNode(), teacher=nn.Identity())
        model = AddLayers()
        dynamic_mixin(
            model,
            DomainAdaptation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.train()
        input = {"real": torch.randn(1), "synthetic": torch.randn(1)}
        output = model(input)
        self.assertEqual(set(output.keys()), {"real", "synthetic", "add"})
        self.assertTrue(torch.allclose(output["real"], (input["real"] + 3.0) * 0.1))
        self.assertTrue(torch.allclose(
            output["synthetic"], (input["synthetic"] + 3.0) * 0.5
        ))
        self.assertTrue(torch.allclose(
            output["add"], ((input["real"] + 1.0) + (input["synthetic"] + 1.0)) * 10.0
        ))
//  [python]
def test_da_train(self):
        """Check train pass results in updated loss output"""
        distillation_helper = TestDAHelper(cfg=CfgNode(), teacher=nn.Identity())
        model = AddLayers()
        dynamic_mixin(
            model,
            DomainAdaptation,
            init_dict={"distillation_helper": distillation_helper},
        )
        model.train()
        input = {"real": torch.randn(1), "synthetic": torch.randn(1)}
        output = model(input)
        self.assertEqual(set(output.keys()), {"real", "synthetic", "add"})
        torch.testing.assert_close(output["real"], (input["real"] + 3.0) * 0.1)
        torch.testing.assert_close(
            output["synthetic"], (input["synthetic"] + 3.0) * 0.5
        )
        torch.testing.assert_close(
            output["add"], ((input["real"] + 1.0) + (input["synthetic"] + 1.0)) * 10.0
        )
//  [python]
def test_default_loss_combiner(self):
        """Check combiner multiplies loss by weights"""
        weights = {"a": torch.randn(1), "b": torch.randn(1)}
        combiner = DefaultLossCombiner(weights)
        input = {"a": 1.0, "b": 10.0}
        output = combiner(input)
        self.assertEqual(output["a"], input["a"] * weights["a"])
        self.assertEqual(output["b"], input["b"] * weights["b"])
//  [python]
def test_default_loss_combiner(self):
        """Check combiner multiplies loss by weights"""
        weights = {"a": torch.randn(1), "b": torch.randn(1)}
        combiner = DefaultLossCombiner(weights)
        input = {"a": 1.0, "b": 10.0}
        output = combiner(input)
        torch.testing.assert_close(output["a"], input["a"] * weights["a"])
        torch.testing.assert_close(output["b"], input["b"] * weights["b"])
//  [python]
def test_unapply(self):
        """Check removing distillation"""
        model = Noop()
        cfg = _get_default_cfg()
        with mock.patch(self._build_teacher_ref):
            hook = DistillationModelingHook(cfg)
        hook.apply(model)
        hook.unapply(model)

        for distillation_attr in [
            "distillation_helper",
            "_original_model_class",
        ]:
            self.assertFalse(hasattr(model, distillation_attr))

        # check forward is the original noop
        batched_inputs = _get_input_data(n=2)
        gt = batched_inputs.detach().clone()
        model.train()
        output = model(batched_inputs)
        t.assert_close(output, gt)
//  [python]
def test_unapply(self):
        """Check removing distillation"""
        model = Noop()
        cfg = _get_default_cfg()
        with mock.patch(self._build_teacher_ref):
            hook = DistillationModelingHook(cfg)
        hook.apply(model)
        hook.unapply(model)

        for distillation_attr in [
            "distillation_helper",
            "_original_model_class",
        ]:
            self.assertFalse(hasattr(model, distillation_attr))

        # check forward is the original noop
        batched_inputs = _get_input_data(n=2)
        gt = batched_inputs.detach().clone()
        model.train()
        output = model(batched_inputs)
        torch.testing.assert_close(output, gt)
//  [python]
def test_apply(self):
        """Check new model has distillation methods"""
        model = Noop()
        model.test_attr = "12345"
        cfg = _get_default_cfg()
        cfg.DISTILLATION.HELPER = "TestHelper"
        with mock.patch(self._build_teacher_ref):
            hook = DistillationModelingHook(cfg)
        hook.apply(model)
        # set teacher manually to override _build_teacher
        model.pseudo_labeler.teacher = DivideInputBy2()

        # check distillation attrs
        self.assertTrue(isinstance(model.distillation_helper, TestHelper))
        self.assertEqual(model._original_model_class, Noop)

        # check retains attrs
        self.assertTrue(hasattr(model, "test_attr"))
        self.assertEqual(model.test_attr, "12345")

        # check inference uses the baseline model which is a noop
        batched_inputs = _get_input_data(n=2)
        model.eval()
        gt = batched_inputs.detach().clone()
        output = model(batched_inputs)
        self.assertTrue(torch.allclose(output, gt))

        # check training uses the pseudo labeler
        model.train()
        gt = [x / 2.0 for x in batched_inputs]
        output = model(batched_inputs)
        self.assertTrue(torch.allclose(output, gt))
//  [python]
def test_apply(self):
        """Check new model has distillation methods"""
        model = Noop()
        model.test_attr = "12345"
        cfg = _get_default_cfg()
        cfg.DISTILLATION.HELPER = "TestHelper"
        with mock.patch(self._build_teacher_ref):
            hook = DistillationModelingHook(cfg)
        hook.apply(model)
        # set teacher manually to override _build_teacher
        model.pseudo_labeler.teacher = DivideInputBy2()

        # check distillation attrs
        self.assertTrue(isinstance(model.distillation_helper, TestHelper))
        self.assertEqual(model._original_model_class, Noop)

        # check retains attrs
        self.assertTrue(hasattr(model, "test_attr"))
        self.assertEqual(model.test_attr, "12345")

        # check inference uses the baseline model which is a noop
        batched_inputs = _get_input_data(n=2)
        model.eval()
        gt = batched_inputs.detach().clone()
        output = model(batched_inputs)
        torch.testing.assert_close(output, gt)

        # check training uses the pseudo labeler
        model.train()
        gt = [x / 2.0 for x in batched_inputs]
        output = model(batched_inputs)
        torch.testing.assert_close(output, gt)
//  [python]
def test_cart_home(self):
        x = torch.randn(2, 3, 3, device=self.device)
        y = pp.homo2cart(pp.cart2homo(x))
        torch.testing.assert_allclose(x, y), 'conversion wrong!'
//  [python]
def test_cart_home(self):
        x = torch.randn(2, 3, 3, device=self.device)
        y = pp.homo2cart(pp.cart2homo(x))
        torch.testing.assert_close(x, y), 'conversion wrong!'
//  [python]
def test_tensor_parallel_linear():
    fabric = Mock()
    fabric.world_size = 4
    fabric.global_rank = 2

    def get_linear(bias=True):
        linear = torch.nn.Linear(8, 8, bias=bias)
        linear.weight.data = torch.arange(64, dtype=torch.float32).reshape(8, 8)
        if bias:
            linear.bias.data = torch.arange(8, dtype=torch.float32)
        return linear

    linear = get_linear()
    tensor_parallel_linear(fabric, linear, "colwise")
    expected = torch.arange(32, 48, dtype=torch.float32).reshape(2, 8)
    t.assert_close(linear.weight, expected)
    expected = torch.arange(4, 6, dtype=torch.float32)
    t.assert_close(linear.bias, expected)

    linear = get_linear(bias=False)
    tensor_parallel_linear(fabric, linear, "rowwise")
    expected = torch.arange(4, 62, 8, dtype=torch.float32).reshape(8, 1)
    expected = torch.cat([expected, expected + 1], dim=1)
    t.assert_close(linear.weight, expected)
    assert linear.bias is None
//  [python]
def test_tensor_parallel_linear():
    fabric = Mock()
    fabric.world_size = 4
    fabric.global_rank = 2

    def get_linear(bias=True):
        linear = torch.nn.Linear(8, 8, bias=bias)
        linear.weight.data = torch.arange(64, dtype=torch.float32).reshape(8, 8)
        if bias:
            linear.bias.data = torch.arange(8, dtype=torch.float32)
        return linear

    linear = get_linear()
    tensor_parallel_linear(fabric, linear, "colwise")
    expected = torch.arange(32, 48, dtype=torch.float32).reshape(2, 8)
    torch.testing.assert_close(linear.weight, expected)
    expected = torch.arange(4, 6, dtype=torch.float32)
    torch.testing.assert_close(linear.bias, expected)

    linear = get_linear(bias=False)
    tensor_parallel_linear(fabric, linear, "rowwise")
    expected = torch.arange(4, 62, 8, dtype=torch.float32).reshape(8, 1)
    expected = torch.cat([expected, expected + 1], dim=1)
    torch.testing.assert_close(linear.weight, expected)
    assert linear.bias is None
//  [python]
def test_ensemble():
    ckpts = [join(dirname(dirname(__file__)), "tests", "example.ckpt")] * 3
    model = load_model(ckpts[0])
    ensemble_model = load_model(ckpts, return_std=True)
    z, pos, batch = create_example_batch(n_atoms=5)

    pred, deriv = model(z, pos, batch)
    pred_ensemble, deriv_ensemble, y_std, neg_dy_std = ensemble_model(z, pos, batch)

    t.assert_allclose(pred, pred_ensemble, atol=1e-5, rtol=1e-5)
    t.assert_allclose(deriv, deriv_ensemble, atol=1e-5, rtol=1e-5)
    assert y_std.shape == pred.shape
    assert neg_dy_std.shape == deriv.shape
    assert (y_std == 0).all()
    assert (neg_dy_std == 0).all()

    import zipfile
    import tempfile

    with tempfile.TemporaryDirectory() as tmpdir:
        ensemble_zip = join(tmpdir, "ensemble.zip")
        with zipfile.ZipFile(ensemble_zip, "w") as zipf:
            for i, ckpt in enumerate(ckpts):
                zipf.write(ckpt, f"model_{i}.ckpt")
        ensemble_model = load_model(ensemble_zip, return_std=True)
        pred_ensemble, deriv_ensemble, y_std, neg_dy_std = ensemble_model(z, pos, batch)

    t.assert_allclose(pred, pred_ensemble, atol=1e-5, rtol=1e-5)
    t.assert_allclose(deriv, deriv_ensemble, atol=1e-5, rtol=1e-5)
    assert y_std.shape == pred.shape
    assert neg_dy_std.shape == deriv.shape
    assert (y_std == 0).all()
    assert (neg_dy_std == 0).all()
//  [python]
def test_ensemble():
    ckpts = [join(dirname(dirname(__file__)), "tests", "example.ckpt")] * 3
    model = load_model(ckpts[0])
    ensemble_model = load_model(ckpts, return_std=True)
    z, pos, batch = create_example_batch(n_atoms=5)

    pred, deriv = model(z, pos, batch)
    pred_ensemble, deriv_ensemble, y_std, neg_dy_std = ensemble_model(z, pos, batch)

    torch.testing.assert_close(pred, pred_ensemble, atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(deriv, deriv_ensemble, atol=1e-5, rtol=1e-5)
    assert y_std.shape == pred.shape
    assert neg_dy_std.shape == deriv.shape
    assert (y_std == 0).all()
    assert (neg_dy_std == 0).all()

    import zipfile
    import tempfile

    with tempfile.TemporaryDirectory() as tmpdir:
        ensemble_zip = join(tmpdir, "ensemble.zip")
        with zipfile.ZipFile(ensemble_zip, "w") as zipf:
            for i, ckpt in enumerate(ckpts):
                zipf.write(ckpt, f"model_{i}.ckpt")
        ensemble_model = load_model(ensemble_zip, return_std=True)
        pred_ensemble, deriv_ensemble, y_std, neg_dy_std = ensemble_model(z, pos, batch)

    torch.testing.assert_close(pred, pred_ensemble, atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(deriv, deriv_ensemble, atol=1e-5, rtol=1e-5)
    assert y_std.shape == pred.shape
    assert neg_dy_std.shape == deriv.shape
    assert (y_std == 0).all()
    assert (neg_dy_std == 0).all()
//  [python]
def test_forward_output(model_name, output_model, overwrite_reference=False):
    pl.seed_everything(1234)

    # create model and sample batch
    derivative = output_model in ["Scalar", "EquivariantScalar"]
    args = load_example_args(
        model_name,
        remove_prior=True,
        output_model=output_model,
        derivative=derivative,
    )
    model = create_model(args)
    z, pos, batch = create_example_batch(n_atoms=5)

    # run step
    pred, deriv = model(z, pos, batch)

    # load reference outputs
    expected_path = join(dirname(__file__), "expected.pkl")
    assert exists(expected_path), "Couldn't locate reference outputs."
    with open(expected_path, "rb") as f:
        expected = pickle.load(f)
    if model_name not in expected or output_model not in expected[model_name]:
        raise ValueError(
            "Model not found in reference outputs, consider running this test with overwrite_reference=True."
        )

    if overwrite_reference:
        # this overwrites the previous reference outputs and shouldn't be executed during testing
        if model_name in expected:
            expected[model_name][output_model] = dict(pred=pred, deriv=deriv)
        else:
            expected[model_name] = {output_model: dict(pred=pred, deriv=deriv)}

        with open(expected_path, "wb") as f:
            pickle.dump(expected, f)
        assert (
            False
        ), f"Set new reference outputs for {model_name} with output model {output_model}."

    # compare actual ouput with reference
    t.assert_allclose(
        pred, expected[model_name][output_model]["pred"], atol=1e-5, rtol=1e-5
    )
    if derivative:
        t.assert_allclose(
            deriv, expected[model_name][output_model]["deriv"], atol=1e-5, rtol=1e-5
        )
//  [python]
def test_forward_output(model_name, output_model, overwrite_reference=False):
    pl.seed_everything(1234)

    # create model and sample batch
    derivative = output_model in ["Scalar", "EquivariantScalar"]
    args = load_example_args(
        model_name,
        remove_prior=True,
        output_model=output_model,
        derivative=derivative,
    )
    model = create_model(args)
    z, pos, batch = create_example_batch(n_atoms=5)

    # run step
    pred, deriv = model(z, pos, batch)

    # load reference outputs
    expected_path = join(dirname(__file__), "expected.pkl")
    assert exists(expected_path), "Couldn't locate reference outputs."
    with open(expected_path, "rb") as f:
        expected = pickle.load(f)
    if model_name not in expected or output_model not in expected[model_name]:
        raise ValueError(
            "Model not found in reference outputs, consider running this test with overwrite_reference=True."
        )

    if overwrite_reference:
        # this overwrites the previous reference outputs and shouldn't be executed during testing
        if model_name in expected:
            expected[model_name][output_model] = dict(pred=pred, deriv=deriv)
        else:
            expected[model_name] = {output_model: dict(pred=pred, deriv=deriv)}

        with open(expected_path, "wb") as f:
            pickle.dump(expected, f)
        assert (
            False
        ), f"Set new reference outputs for {model_name} with output model {output_model}."

    # compare actual ouput with reference
    torch.testing.assert_close(
        pred, expected[model_name][output_model]["pred"], atol=1e-5, rtol=1e-5
    )
    if derivative:
        torch.testing.assert_close(
            deriv, expected[model_name][output_model]["deriv"], atol=1e-5, rtol=1e-5
        )
//  [python]
def test_to_tensor_list(self):
        _list = [1, 2, 3]
        torch.testing.assert_allclose(to_tensor(_list), torch.tensor(_list))
//  [python]
def test_to_tensor_list(self):
        _list = [1, 2, 3]
        torch.testing.assert_close(to_tensor(_list), torch.tensor(_list))
//  [python]
def test_blur_equivalence(kernel_size: tuple[int, int], sigma: tuple[int, int], channels: int) -> None:
    """Test if the GaussianBlur2d layer produces the same output as Kornia's GaussianBlur2d layer."""
    for _ in range(10):
        input_tensor = torch.randn((3, channels, 128, 128))
        kornia = korniaGaussianBlur2d(kernel_size, sigma, separable=False)
        blur_kornia = kornia(input_tensor)
        gaussian = GaussianBlur2d(kernel_size=kernel_size, sigma=sigma, channels=channels)
        blur_gaussian = gaussian(input_tensor)
        torch.testing.assert_allclose(blur_kornia, blur_gaussian)
//  [python]
def test_blur_equivalence(kernel_size: tuple[int, int], sigma: tuple[int, int], channels: int) -> None:
    """Test if the GaussianBlur2d layer produces the same output as Kornia's GaussianBlur2d layer."""
    for _ in range(10):
        input_tensor = torch.randn((3, channels, 128, 128))
        kornia = korniaGaussianBlur2d(kernel_size, sigma, separable=False)
        blur_kornia = kornia(input_tensor)
        gaussian = GaussianBlur2d(kernel_size=kernel_size, sigma=sigma, channels=channels)
        blur_gaussian = gaussian(input_tensor)
        torch.testing.assert_close(blur_kornia, blur_gaussian)
//  [python]
def test_joint_limits(self):
        # panda_joint_lims:
        # (-2.8973, 2.8973),
        # (-1.7628, 1.7628),
        # (-2.8973, 2.8973),
        # (-3.0718, -0.0698),
        # (-2.8973, 2.8973),
        # (-0.0175, 3.7525),
        # (-2.8973, 2.8973)

        # Test 1:
        joint_angles_unclamped = torch.tensor([
            [0, 0, 0, -0.1, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [3, 0, 0, 0, 0, 0, 0],
        ])
        panda = Panda()
        returned = panda.clamp_to_joint_limits(joint_angles_unclamped)
        expected = torch.tensor([
            [0, 0, 0, -0.1, 0, 0, 0],
            [0, 0, 0, -0.0698, 0, 0, 0],
            [2.8973, 0, 0, -0.0698, 0, 0, 0],
        ])
        torch.testing.assert_allclose(returned, expected)
//  [python]
def test_joint_limits(self):
        # panda_joint_lims:
        # (-2.8973, 2.8973),
        # (-1.7628, 1.7628),
        # (-2.8973, 2.8973),
        # (-3.0718, -0.0698),
        # (-2.8973, 2.8973),
        # (-0.0175, 3.7525),
        # (-2.8973, 2.8973)

        # Test 1:
        joint_angles_unclamped = torch.tensor([
            [0, 0, 0, -0.1, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [3, 0, 0, 0, 0, 0, 0],
        ])
        panda = Panda()
        returned = panda.clamp_to_joint_limits(joint_angles_unclamped)
        expected = torch.tensor([
            [0, 0, 0, -0.1, 0, 0, 0],
            [0, 0, 0, -0.0698, 0, 0, 0],
            [2.8973, 0, 0, -0.0698, 0, 0, 0],
        ])
        torch.testing.assert_close(returned, expected)
//  [python]
def test_model_8b_chat_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = PersimmonForCausalLM.from_pretrained(
            "adept/persimmon-8b-chat", load_in_8bit=True, device_map={"": 0}, torch_dtype=torch.float16
        )
        out = model(torch.tensor([input_ids], device=torch_device)).logits

        EXPECTED_MEAN = torch.tensor(
            [[-11.4726, -11.1495, -11.2694, -11.2223, -10.9452, -11.0663, -11.0031, -11.1028]]
        )
        # change dtype to `torch.float32` before calling `mean` to avoid `nan` values
        t.assert_close(out.cpu().to(torch.float32).mean(-1), EXPECTED_MEAN, atol=1e-4, rtol=1e-4)
        # fmt: off
        EXPECTED_SLICE = torch.tensor(
            [-16.9062, -16.9062, -16.9062, -16.9062, -16.8906, -16.9062, -16.9531, -16.9062, -16.9062, -16.9062, -16.9531, -16.9062, -16.9531, -16.9062, -16.9062, -16.9062, -16.9062, -16.9062, -16.9531, -16.9062, -16.9062, -16.9062, -16.9062, -16.9062, -16.9062, -16.9531, -16.9062, -16.9531, -16.9062, -16.9062],
            dtype=torch.float16
        )
        # fmt: on
        t.assert_close(out.cpu()[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)

        backend_empty_cache(torch_device)
        del model
        gc.collect()
//  [python]
def test_model_8b_chat_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = PersimmonForCausalLM.from_pretrained(
            "adept/persimmon-8b-chat", load_in_8bit=True, device_map={"": 0}, torch_dtype=torch.float16
        )
        out = model(torch.tensor([input_ids], device=torch_device)).logits

        EXPECTED_MEAN = torch.tensor(
            [[-11.4726, -11.1495, -11.2694, -11.2223, -10.9452, -11.0663, -11.0031, -11.1028]]
        )
        # change dtype to `torch.float32` before calling `mean` to avoid `nan` values
        torch.testing.assert_close(out.cpu().to(torch.float32).mean(-1), EXPECTED_MEAN, atol=1e-4, rtol=1e-4)
        # fmt: off
        EXPECTED_SLICE = torch.tensor(
            [-16.9062, -16.9062, -16.9062, -16.9062, -16.8906, -16.9062, -16.9531, -16.9062, -16.9062, -16.9062, -16.9531, -16.9062, -16.9531, -16.9062, -16.9062, -16.9062, -16.9062, -16.9062, -16.9531, -16.9062, -16.9062, -16.9062, -16.9062, -16.9062, -16.9062, -16.9531, -16.9062, -16.9531, -16.9062, -16.9062],
            dtype=torch.float16
        )
        # fmt: on
        torch.testing.assert_close(out.cpu()[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)

        backend_empty_cache(torch_device)
        del model
        gc.collect()
//  [python]
def test_to_tensor_array(self):
        _array = np.array([1, 2, 3])
        torch.testing.assert_allclose(to_tensor(_array), torch.tensor(_array))
//  [python]
def test_to_tensor_array(self):
        _array = np.array([1, 2, 3])
        torch.testing.assert_close(to_tensor(_array), torch.tensor(_array))
//  [python]
def test_to_tensor_sequence_(self):
        _sequence = [[1.0, 2.0, 3.0]]
        torch.testing.assert_allclose(to_tensor(_sequence), torch.tensor(_sequence))
//  [python]
def test_to_tensor_sequence_(self):
        _sequence = [[1.0, 2.0, 3.0]]
        torch.testing.assert_close(to_tensor(_sequence), torch.tensor(_sequence))
//  [python]
onnxruntime_outputs = ort_session.run(None, onnxruntime_input)

torch_outputs = aten_add_model(input_add_x, input_add_y)
torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)

assert len(torch_outputs) == len(onnxruntime_outputs)
for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
    torch.testing.assert_allclose(torch_output, torch.tensor(onnxruntime_output))
//  [python]
onnxruntime_outputs = ort_session.run(None, onnxruntime_input)

torch_outputs = aten_add_model(input_add_x, input_add_y)
torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)

assert len(torch_outputs) == len(onnxruntime_outputs)
for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
    torch.testing.assert_close(torch_output, torch.tensor(onnxruntime_output))
//  [python]
def test_inference_mask_generation_two_points_point_batch(self):
        model = SamModel.from_pretrained("facebook/sam-vit-base")
        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

        model.to(torch_device)
        model.eval()

        raw_image = prepare_image()

        input_points = torch.Tensor([[[400, 650]], [[220, 470]]]).cpu()  # fmt: skip

        input_points = input_points.unsqueeze(0)

        inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        iou_scores = outputs.iou_scores.cpu()
        self.assertTrue(iou_scores.shape == (1, 2, 3))
        torch.allclose(
            iou_scores, torch.tensor([[[0.9105, 0.9825, 0.9675], [0.7646, 0.7943, 0.7774]]]), atol=1e-4, rtol=1e-4
        )
//  [python]
def test_split_configs_to_revolute_and_prismatic(self):
        """Test that split_configs_to_revolute_and_prismatic() returns as expected"""
        # Test 1: Fetch
        fetch = Fetch()
        qs = torch.tensor(
            [
                [0.1, 0.2, 0.3, 0.5, 0.6, 0.7, 0.8, 0.9],
                [1.1, 1.2, 1.3, 1.5, 1.6, 1.7, 1.8, 1.9],
                [2.1, 2.2, 2.3, 2.5, 2.6, 2.7, 2.8, 2.9],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_revolute = torch.tensor(
            [
                [0.2, 0.3, 0.5, 0.6, 0.7, 0.8, 0.9],
                [1.2, 1.3, 1.5, 1.6, 1.7, 1.8, 1.9],
                [2.2, 2.3, 2.5, 2.6, 2.7, 2.8, 2.9],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_prismatic = torch.tensor(
            [
                [0.1],
                [1.1],
                [2.1],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        revolute, prismatic = fetch.split_configs_to_revolute_and_prismatic(qs)
        torch.testing.assert_allclose(expected_prismatic, prismatic)
        torch.testing.assert_allclose(expected_revolute, revolute)

        # Test 2: FetchArm
        fetch = FetchArm()
        qs = torch.tensor(
            [
                [0.1, 0.2, 0.3, 0.5, 0.6, 0.7, 0.8],
                [1.1, 1.2, 1.3, 1.5, 1.6, 1.7, 1.8],
                [2.1, 2.2, 2.3, 2.5, 2.6, 2.7, 2.8],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_revolute = torch.tensor(
            [
                [0.1, 0.2, 0.3, 0.5, 0.6, 0.7, 0.8],
                [1.1, 1.2, 1.3, 1.5, 1.6, 1.7, 1.8],
                [2.1, 2.2, 2.3, 2.5, 2.6, 2.7, 2.8],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_prismatic = torch.tensor(
            [
                [],
                [],
                [],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        revolute, prismatic = fetch.split_configs_to_revolute_and_prismatic(qs)
        torch.testing.assert_allclose(expected_prismatic, prismatic)
        torch.testing.assert_allclose(expected_revolute, revolute)
//  [python]
def test_split_configs_to_revolute_and_prismatic(self):
        """Test that split_configs_to_revolute_and_prismatic() returns as expected"""
        # Test 1: Fetch
        fetch = Fetch()
        qs = torch.tensor(
            [
                [0.1, 0.2, 0.3, 0.5, 0.6, 0.7, 0.8, 0.9],
                [1.1, 1.2, 1.3, 1.5, 1.6, 1.7, 1.8, 1.9],
                [2.1, 2.2, 2.3, 2.5, 2.6, 2.7, 2.8, 2.9],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_revolute = torch.tensor(
            [
                [0.2, 0.3, 0.5, 0.6, 0.7, 0.8, 0.9],
                [1.2, 1.3, 1.5, 1.6, 1.7, 1.8, 1.9],
                [2.2, 2.3, 2.5, 2.6, 2.7, 2.8, 2.9],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_prismatic = torch.tensor(
            [
                [0.1],
                [1.1],
                [2.1],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        revolute, prismatic = fetch.split_configs_to_revolute_and_prismatic(qs)
        torch.testing.assert_close(expected_prismatic, prismatic)
        torch.testing.assert_close(expected_revolute, revolute)

        # Test 2: FetchArm
        fetch = FetchArm()
        qs = torch.tensor(
            [
                [0.1, 0.2, 0.3, 0.5, 0.6, 0.7, 0.8],
                [1.1, 1.2, 1.3, 1.5, 1.6, 1.7, 1.8],
                [2.1, 2.2, 2.3, 2.5, 2.6, 2.7, 2.8],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_revolute = torch.tensor(
            [
                [0.1, 0.2, 0.3, 0.5, 0.6, 0.7, 0.8],
                [1.1, 1.2, 1.3, 1.5, 1.6, 1.7, 1.8],
                [2.1, 2.2, 2.3, 2.5, 2.6, 2.7, 2.8],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        expected_prismatic = torch.tensor(
            [
                [],
                [],
                [],
            ],
            dtype=torch.float32,
            device="cpu",
        )
        revolute, prismatic = fetch.split_configs_to_revolute_and_prismatic(qs)
        torch.testing.assert_close(expected_prismatic, prismatic)
        torch.testing.assert_close(expected_revolute, revolute)
//  [python]
def test_inference_mask_generation_three_boxes_point_batch(self):
        model = SamModel.from_pretrained("facebook/sam-vit-base")
        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

        model.to(torch_device)
        model.eval()

        raw_image = prepare_image()

        # fmt: off
        input_boxes = torch.Tensor([[[620, 900, 1000, 1255]], [[75, 275, 1725, 850]],  [[75, 275, 1725, 850]]]).cpu()
        EXPECTED_IOU = torch.tensor([[[0.9773, 0.9881, 0.9522],
         [0.5996, 0.7661, 0.7937],
         [0.5996, 0.7661, 0.7937]]])
        # fmt: on
        input_boxes = input_boxes.unsqueeze(0)

        inputs = processor(raw_image, input_boxes=input_boxes, return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        iou_scores = outputs.iou_scores.cpu()
        self.assertTrue(iou_scores.shape == (1, 3, 3))
        torch.allclose(iou_scores, EXPECTED_IOU, atol=1e-4, rtol=1e-4)