jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (7,), (2,)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        query = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16),
                torch.randn(2, 8, 16, 32),
                torch.randn(4, 1, 32, 16),
            ],
            **kwargs,
        )
        key = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16),
                torch.randn(2, 8, 16, 32),
                torch.randn(4, 1, 32, 16),
            ],
            **kwargs,
        )
        out_nested = na1d_qk(query, key, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for q, k in zip(query, key):
            out_ref.append(na1d_qk(q, k, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            t.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            additional_keys = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na1d_qk(
                query,
                key,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_keys=additional_keys,
            )
            out_ref = []
            for q, k, k_add in zip(query, key, additional_keys):
                out_ref.append(
                    na1d_qk(
                        q,
                        k,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_keys=k_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                t.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (7,), (2,)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        query = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16),
                torch.randn(2, 8, 16, 32),
                torch.randn(4, 1, 32, 16),
            ],
            **kwargs,
        )
        key = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16),
                torch.randn(2, 8, 16, 32),
                torch.randn(4, 1, 32, 16),
            ],
            **kwargs,
        )
        out_nested = na1d_qk(query, key, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for q, k in zip(query, key):
            out_ref.append(na1d_qk(q, k, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            additional_keys = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na1d_qk(
                query,
                key,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_keys=additional_keys,
            )
            out_ref = []
            for q, k, k_add in zip(query, key, additional_keys):
                out_ref.append(
                    na1d_qk(
                        q,
                        k,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_keys=k_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def test_convs(devices, sharded, extra_options):
    torch.manual_seed(0)

    batchnorm_cls = (None, nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)
    # ^-- note: batchnorms test that tensor_parallel handles buffers (non-parameter state tensors) correctly
    for Conv, nd in (
        (nn.Conv1d, 1),
        (nn.Conv2d, 2),
        (nn.Conv3d, 3),
        (nn.ConvTranspose1d, 1),
        (nn.ConvTranspose2d, 2),
        (nn.ConvTranspose3d, 3),
    ):
        if issubclass(Conv, _ConvTransposeNd) and "padding" in extra_options:
            continue  # unsupported by pytorch
        model = nn.Sequential(
            Conv(32, 64, kernel_size=(3,) * nd, **extra_options),
            batchnorm_cls[nd](64),
            nn.ReLU(),
            Conv(64, 32, kernel_size=(3,) * nd, **extra_options),
        )
        inputs1 = torch.randn(3, 32, *[10 for _ in range(nd)], requires_grad=True)
        inputs2 = inputs1.detach().clone().requires_grad_(True)
        ref_out = model(inputs1)
        ref_out.norm().backward()

        model_tp = deepcopy(model)  # deepcopy to avoid accidental grad spillage and false positives
        model_tp = TensorParallel(model_tp, device_ids=devices, sharded=sharded)
        out_ours = model_tp(inputs2)
        out_ours.norm().backward()
        t.assert_close(
            ref_out,
            out_ours,
            atol=1e-6,
            rtol=1e-3,
            msg=lambda msg: f"{msg}\n where Conv is {Conv} with extra_options {extra_options}",
        )
        t.assert_close(
            inputs1.grad,
            inputs2.grad,
            atol=1e-3,
            rtol=1e-05,
            msg=lambda msg: f"{msg}\n where Conv is {Conv} with extra_options {extra_options}",
        )
//  [python]
def test_convs(devices, sharded, extra_options):
    torch.manual_seed(0)

    batchnorm_cls = (None, nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)
    # ^-- note: batchnorms test that tensor_parallel handles buffers (non-parameter state tensors) correctly
    for Conv, nd in (
        (nn.Conv1d, 1),
        (nn.Conv2d, 2),
        (nn.Conv3d, 3),
        (nn.ConvTranspose1d, 1),
        (nn.ConvTranspose2d, 2),
        (nn.ConvTranspose3d, 3),
    ):
        if issubclass(Conv, _ConvTransposeNd) and "padding" in extra_options:
            continue  # unsupported by pytorch
        model = nn.Sequential(
            Conv(32, 64, kernel_size=(3,) * nd, **extra_options),
            batchnorm_cls[nd](64),
            nn.ReLU(),
            Conv(64, 32, kernel_size=(3,) * nd, **extra_options),
        )
        inputs1 = torch.randn(3, 32, *[10 for _ in range(nd)], requires_grad=True)
        inputs2 = inputs1.detach().clone().requires_grad_(True)
        ref_out = model(inputs1)
        ref_out.norm().backward()

        model_tp = deepcopy(model)  # deepcopy to avoid accidental grad spillage and false positives
        model_tp = TensorParallel(model_tp, device_ids=devices, sharded=sharded)
        out_ours = model_tp(inputs2)
        out_ours.norm().backward()
        torch.testing.assert_close(
            ref_out,
            out_ours,
            atol=1e-6,
            rtol=1e-3,
            msg=lambda msg: f"{msg}\n where Conv is {Conv} with extra_options {extra_options}",
        )
        torch.testing.assert_close(
            inputs1.grad,
            inputs2.grad,
            atol=1e-3,
            rtol=1e-05,
            msg=lambda msg: f"{msg}\n where Conv is {Conv} with extra_options {extra_options}",
        )
//  [python]
def test_transform_module_registration_buffers() -> None:
    # Check that z is in buffers
    offset = torch.randn((5, 1))
    transform = _OffsetTransformModule(offset)
    buffers = dict(transform.named_buffers())
    t.assert_close(buffers["offset"], offset)
//  [python]
def test_transform_module_registration_buffers() -> None:
    # Check that z is in buffers
    offset = torch.randn((5, 1))
    transform = _OffsetTransformModule(offset)
    buffers = dict(transform.named_buffers())
    torch.testing.assert_close(buffers["offset"], offset)
//  [python]
def test_config_backward_compatibility(devices):
    class Module(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(10, 10, bias=False)

        def forward(self, x):
            return self.linear(x)

    model = Module()

    modern_state_dict = TensorParallel(
        model,
        device_ids=devices,
        tensor_parallel_config=Config(
            state_rules={r".*linear.weight": tensor_parallel.state_actions.Split(world_size=len(devices), dim=0)},
            input_rules={},
            output_rules={},
            attr_rules={},
        ),
    ).state_dict()

    legacy_callable_state_dict = TensorParallel(
        model,
        device_ids=devices,
        tensor_parallel_config=Config(
            state_rules={
                r".*linear.weight": lambda tensor, rank: torch.tensor_split(tensor, len(devices), dim=0)[rank]
            },
            input_rules={},
            output_rules={},
            attr_rules={},
        ),
    ).state_dict()

    legacy_tuple_state_dict = TensorParallel(
        model,
        device_ids=devices,
        tensor_parallel_config=Config(
            state_rules={
                r".*linear.weight": (
                    lambda tensor, rank: torch.tensor_split(tensor, len(devices), dim=0)[rank],
                    lambda tensors: torch.cat([tensor.cpu() for tensor in tensors], dim=0),
                )
            },
            input_rules={},
            output_rules={},
            attr_rules={},
        ),
    ).state_dict()

    for name, item in modern_state_dict.items():
        assert name in legacy_callable_state_dict, f"{name} not in legacy_callable_state_dict"
        t.assert_close(item, legacy_callable_state_dict[name])

        assert name in legacy_tuple_state_dict, f"{name} not in legacy_tuple_state_dict"
        t.assert_close(item, legacy_tuple_state_dict[name])
//  [python]
def test_config_backward_compatibility(devices):
    class Module(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(10, 10, bias=False)

        def forward(self, x):
            return self.linear(x)

    model = Module()

    modern_state_dict = TensorParallel(
        model,
        device_ids=devices,
        tensor_parallel_config=Config(
            state_rules={r".*linear.weight": tensor_parallel.state_actions.Split(world_size=len(devices), dim=0)},
            input_rules={},
            output_rules={},
            attr_rules={},
        ),
    ).state_dict()

    legacy_callable_state_dict = TensorParallel(
        model,
        device_ids=devices,
        tensor_parallel_config=Config(
            state_rules={
                r".*linear.weight": lambda tensor, rank: torch.tensor_split(tensor, len(devices), dim=0)[rank]
            },
            input_rules={},
            output_rules={},
            attr_rules={},
        ),
    ).state_dict()

    legacy_tuple_state_dict = TensorParallel(
        model,
        device_ids=devices,
        tensor_parallel_config=Config(
            state_rules={
                r".*linear.weight": (
                    lambda tensor, rank: torch.tensor_split(tensor, len(devices), dim=0)[rank],
                    lambda tensors: torch.cat([tensor.cpu() for tensor in tensors], dim=0),
                )
            },
            input_rules={},
            output_rules={},
            attr_rules={},
        ),
    ).state_dict()

    for name, item in modern_state_dict.items():
        assert name in legacy_callable_state_dict, f"{name} not in legacy_callable_state_dict"
        torch.testing.assert_close(item, legacy_callable_state_dict[name])

        assert name in legacy_tuple_state_dict, f"{name} not in legacy_tuple_state_dict"
        torch.testing.assert_close(item, legacy_tuple_state_dict[name])
//  [python]
def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (7,), (2,)
        num_na_weights = get_num_na_weights(kernel_size)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        attn = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, num_na_weights),
                torch.randn(2, 8, 16, num_na_weights),
                torch.randn(4, 1, 32, num_na_weights),
            ],
            **kwargs,
        )
        value = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16),
                torch.randn(2, 8, 16, 32),
                torch.randn(4, 1, 32, 16),
            ],
            **kwargs,
        )
        out_nested = na1d_av(attn, value, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for a, v in zip(attn, value):
            out_ref.append(na1d_av(a, v, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            attn = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, 14, num_na_weights + n_extra_tokens),
                    torch.randn(2, 8, 16, num_na_weights + n_extra_tokens),
                    torch.randn(4, 1, 32, num_na_weights + n_extra_tokens),
                ],
                **kwargs,
            )
            additional_values = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na1d_av(
                attn,
                value,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_values=additional_values,
            )
            out_ref = []
            for a, v, v_add in zip(attn, value, additional_values):
                out_ref.append(
                    na1d_av(
                        a,
                        v,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_values=v_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (7,), (2,)
        num_na_weights = get_num_na_weights(kernel_size)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        attn = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, num_na_weights),
                torch.randn(2, 8, 16, num_na_weights),
                torch.randn(4, 1, 32, num_na_weights),
            ],
            **kwargs,
        )
        value = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16),
                torch.randn(2, 8, 16, 32),
                torch.randn(4, 1, 32, 16),
            ],
            **kwargs,
        )
        out_nested = na1d_av(attn, value, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for a, v in zip(attn, value):
            out_ref.append(na1d_av(a, v, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            attn = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, 14, num_na_weights + n_extra_tokens),
                    torch.randn(2, 8, 16, num_na_weights + n_extra_tokens),
                    torch.randn(4, 1, 32, num_na_weights + n_extra_tokens),
                ],
                **kwargs,
            )
            additional_values = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na1d_av(
                attn,
                value,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_values=additional_values,
            )
            out_ref = []
            for a, v, v_add in zip(attn, value, additional_values):
                out_ref.append(
                    na1d_av(
                        a,
                        v,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_values=v_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def test_sequential_transform_module_inner_buffers() -> None:
    offset = torch.randn((5, 1))
    transform = _OffsetTransformModule(offset)
    seq_transform = SequentialTransformModule[torch.Tensor, torch.Tensor](transform, transform.inv)
    # Check that buffers are stored for the inner transformation
    seq_buffers = dict(seq_transform.named_buffers())
    for name, buffer in transform.named_buffers():
        t.assert_close(seq_buffers[f"0.{name}"], buffer)
//  [python]
def test_sequential_transform_module_inner_buffers() -> None:
    offset = torch.randn((5, 1))
    transform = _OffsetTransformModule(offset)
    seq_transform = SequentialTransformModule[torch.Tensor, torch.Tensor](transform, transform.inv)
    # Check that buffers are stored for the inner transformation
    seq_buffers = dict(seq_transform.named_buffers())
    for name, buffer in transform.named_buffers():
        torch.testing.assert_close(seq_buffers[f"0.{name}"], buffer)
//  [python]
def test_joint_transform_module_inner_buffers() -> None:
    offset = torch.randn((2,))
    transform = _OffsetTransformModule(offset)
    joint_transform = JointTransformModule({"a": transform})
    # Check that buffers are stored for the inner transformation
    joint_buffers = dict(joint_transform.named_buffers())
    for name, buffer in transform.named_buffers():
        t.assert_close(joint_buffers[f"transformations.a.{name}"], buffer)
//  [python]
def test_joint_transform_module_inner_buffers() -> None:
    offset = torch.randn((2,))
    transform = _OffsetTransformModule(offset)
    joint_transform = JointTransformModule({"a": transform})
    # Check that buffers are stored for the inner transformation
    joint_buffers = dict(joint_transform.named_buffers())
    for name, buffer in transform.named_buffers():
        torch.testing.assert_close(joint_buffers[f"transformations.a.{name}"], buffer)
//  [python]
def _test_model_numerically(
        self,
        model,
        dynamo_backend,
        example_args_collection,
        fullgraph: bool = False,
        test_backward: bool = False,
        atol: float = 1e-5,
        rtol: float = 1e-6,
    ):
        """Run original and compiled model and compare the results.

        Args:
            model: The model to test.
            dynamo_backend: The dynamo backend to use. Here we use string `onnxrt` or
              the first returned value of `make_aot_ort(dynamic=True)`.
            example_args_collection: A tuple of example arguments to test. E.g.,
                (
                  (torch.randn(2), torch.randn(2)),
                  (torch.randn(4), torch.randn(4)),
                )
              if you want to test
                model(torch.randn(2), torch.randn(2)) and
                model(torch.randn(4), torch.randn(4))
              .
        """
        compiled_model = torch.compile(
            model if not isinstance(model, torch.nn.Module) else copy.deepcopy(model),
            backend=dynamo_backend,
            dynamic=True,
            fullgraph=fullgraph,
        )

        for example_args in example_args_collection:
            baseline_result = model(*example_args)
            result = compiled_model(*example_args)
            if isinstance(baseline_result, torch.Tensor):
                torch.testing.assert_allclose(
                    baseline_result, result, atol=atol, rtol=rtol
                )
                if test_backward:
                    baseline_result.sum().backward()
                    result.sum().backward()
                    for baseline_param, param in zip(
                        model.parameters(), compiled_model.parameters()
                    ):
                        torch.testing.assert_allclose(
                            baseline_param.grad, param.grad, atol=atol, rtol=rtol
                        )
            else:
                assert (
                    test_backward is False
                ), "Calculating backward with multiple outputs is not supported yet."
                for baseline_elem, result_elem in zip(baseline_result, result):
                    torch.testing.assert_allclose(
                        baseline_elem, result_elem, atol=atol, rtol=rtol
                    )
//  [python]
def _test_model_numerically(
        self,
        model,
        dynamo_backend,
        example_args_collection,
        fullgraph: bool = False,
        test_backward: bool = False,
        atol: float = 1e-5,
        rtol: float = 1e-6,
    ):
        """Run original and compiled model and compare the results.

        Args:
            model: The model to test.
            dynamo_backend: The dynamo backend to use. Here we use string `onnxrt` or
              the first returned value of `make_aot_ort(dynamic=True)`.
            example_args_collection: A tuple of example arguments to test. E.g.,
                (
                  (torch.randn(2), torch.randn(2)),
                  (torch.randn(4), torch.randn(4)),
                )
              if you want to test
                model(torch.randn(2), torch.randn(2)) and
                model(torch.randn(4), torch.randn(4))
              .
        """
        compiled_model = torch.compile(
            model if not isinstance(model, torch.nn.Module) else copy.deepcopy(model),
            backend=dynamo_backend,
            dynamic=True,
            fullgraph=fullgraph,
        )

        for example_args in example_args_collection:
            baseline_result = model(*example_args)
            result = compiled_model(*example_args)
            if isinstance(baseline_result, torch.Tensor):
                torch.testing.assert_close(
                    baseline_result, result, atol=atol, rtol=rtol
                )
                if test_backward:
                    baseline_result.sum().backward()
                    result.sum().backward()
                    for baseline_param, param in zip(
                        model.parameters(), compiled_model.parameters()
                    ):
                        torch.testing.assert_close(
                            baseline_param.grad, param.grad, atol=atol, rtol=rtol
                        )
            else:
                assert (
                    test_backward is False
                ), "Calculating backward with multiple outputs is not supported yet."
                for baseline_elem, result_elem in zip(baseline_result, result):
                    torch.testing.assert_close(
                        baseline_elem, result_elem, atol=atol, rtol=rtol
                    )
//  [python]
def test_add_state_tensor(self) -> None:
        metric = DummySumMetric()
        metric._add_state("x1", torch.tensor([0.0, 1.0]))
        t.assert_close(metric.x1, torch.tensor([0.0, 1.0]))

        default_tensor = torch.tensor([0.0, 1.0])
        metric._add_state("x2", default_tensor)
        t.assert_close(metric.x2, default_tensor)
        # state is still set to the value when _add_state() is called
        default_tensor = torch.tensor([1.0, 2.0])
        self.assertFalse(torch.allclose(metric.x2, default_tensor))
        t.assert_close(metric.x2, torch.tensor([0.0, 1.0]))
//  [python]
def test_add_state_tensor(self) -> None:
        metric = DummySumMetric()
        metric._add_state("x1", torch.tensor([0.0, 1.0]))
        torch.testing.assert_close(metric.x1, torch.tensor([0.0, 1.0]))

        default_tensor = torch.tensor([0.0, 1.0])
        metric._add_state("x2", default_tensor)
        torch.testing.assert_close(metric.x2, default_tensor)
        # state is still set to the value when _add_state() is called
        default_tensor = torch.tensor([1.0, 2.0])
        self.assertFalse(torch.allclose(metric.x2, default_tensor))
        torch.testing.assert_close(metric.x2, torch.tensor([0.0, 1.0]))
//  [python]
def test_add_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric()
        metric._add_state("x1", [])
        t.assert_close(metric.x1, [])

        metric._add_state("x2", [torch.tensor(0.0), torch.tensor(1.0)])
        t.assert_close(metric.x2, [torch.tensor(0.0), torch.tensor(1.0)])

        tensor_list = [torch.tensor(0.0)]
        metric._add_state("x3", tensor_list)
        t.assert_close(metric.x3, tensor_list)
        tensor_list.append(torch.tensor(1.0))
        self.assertNotEqual(metric.x3, tensor_list)
        t.assert_close(metric.x3, [torch.tensor(0.0)])
//  [python]
def test_add_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric()
        metric._add_state("x1", [])
        torch.testing.assert_close(metric.x1, [])

        metric._add_state("x2", [torch.tensor(0.0), torch.tensor(1.0)])
        torch.testing.assert_close(metric.x2, [torch.tensor(0.0), torch.tensor(1.0)])

        tensor_list = [torch.tensor(0.0)]
        metric._add_state("x3", tensor_list)
        torch.testing.assert_close(metric.x3, tensor_list)
        tensor_list.append(torch.tensor(1.0))
        self.assertNotEqual(metric.x3, tensor_list)
        torch.testing.assert_close(metric.x3, [torch.tensor(0.0)])
//  [python]
def test_registration(data: X, transform: TransformModule[X, Y], _, to_kwargs: dict[str, Any]) -> None:
    """Test that registration is working by testing that the state can be moved and loaded."""
    transform_modified: TransformModule[X, Y] = transform.to(**to_kwargs)

    # Collect parameters and buffers as tensors
    tensors = dict(itertools.chain(transform.named_buffers(), transform.named_parameters()))
    tensors_modified = dict(itertools.chain(transform_modified.named_buffers(), transform_modified.named_parameters()))

    # Check that the tensors are equivalent
    assert set(tensors) == set(tensors_modified)
    for name in tensors:
        t.assert_close(tensors[name].to(**to_kwargs), tensors_modified[name])

    # Check that the state dict is consistent and picklable
    state_dict = transform_modified.state_dict()
    with io.BytesIO() as f:
        torch.save(state_dict, f)
        f.seek(0)
        state_dict = torch.load(f)
    for name in tensors:
        t.assert_close(tensors[name].to(**to_kwargs), state_dict[name])

    # Produce the output for x
    if isinstance(data, TensorDictBase):
        x_modified = data.apply(lambda x_: x_.to(**to_kwargs))
    else:
        x_modified = data.to(**to_kwargs)
    y_modified = transform_modified(x_modified)
    y = transform(data)

    # Check that the output remains correct, i.e. the transformation is approx equivariant w.r.t. the `to` operator.
    if isinstance(y, TensorDictBase):
        assert isinstance(y_modified, TensorDictBase)  # plays nicer with mypy than checking type equality
        for key in y.keys():
            t.assert_close(y_modified.get(key), y.get(key).to(**to_kwargs), atol=2e-2, rtol=1e-2)
    else:
        assert isinstance(y_modified, torch.Tensor)  # plays nicer with mypy than checking type equality
        t.assert_close(y_modified, y.to(**to_kwargs), atol=2e-2, rtol=1e-2)
//  [python]
def test_registration(data: X, transform: TransformModule[X, Y], _, to_kwargs: dict[str, Any]) -> None:
    """Test that registration is working by testing that the state can be moved and loaded."""
    transform_modified: TransformModule[X, Y] = transform.to(**to_kwargs)

    # Collect parameters and buffers as tensors
    tensors = dict(itertools.chain(transform.named_buffers(), transform.named_parameters()))
    tensors_modified = dict(itertools.chain(transform_modified.named_buffers(), transform_modified.named_parameters()))

    # Check that the tensors are equivalent
    assert set(tensors) == set(tensors_modified)
    for name in tensors:
        torch.testing.assert_close(tensors[name].to(**to_kwargs), tensors_modified[name])

    # Check that the state dict is consistent and picklable
    state_dict = transform_modified.state_dict()
    with io.BytesIO() as f:
        torch.save(state_dict, f)
        f.seek(0)
        state_dict = torch.load(f)
    for name in tensors:
        torch.testing.assert_close(tensors[name].to(**to_kwargs), state_dict[name])

    # Produce the output for x
    if isinstance(data, TensorDictBase):
        x_modified = data.apply(lambda x_: x_.to(**to_kwargs))
    else:
        x_modified = data.to(**to_kwargs)
    y_modified = transform_modified(x_modified)
    y = transform(data)

    # Check that the output remains correct, i.e. the transformation is approx equivariant w.r.t. the `to` operator.
    if isinstance(y, TensorDictBase):
        assert isinstance(y_modified, TensorDictBase)  # plays nicer with mypy than checking type equality
        for key in y.keys():
            torch.testing.assert_close(y_modified.get(key), y.get(key).to(**to_kwargs), atol=2e-2, rtol=1e-2)
    else:
        assert isinstance(y_modified, torch.Tensor)  # plays nicer with mypy than checking type equality
        torch.testing.assert_close(y_modified, y.to(**to_kwargs), atol=2e-2, rtol=1e-2)
//  [python]
def test_reset_state_tensor(self) -> None:
        metric = DummySumMetric()
        t.assert_close(metric.sum, torch.tensor(0.0))
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        t.assert_close(metric.sum, torch.tensor(3.0))

        # reset once
        metric.reset()
        t.assert_close(metric.sum, torch.tensor(0.0))
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        t.assert_close(metric.sum, torch.tensor(3.0))

        # reset again
        metric.reset()
        t.assert_close(metric.sum, torch.tensor(0.0))

        # update and reset chained
        metric.update(torch.tensor(1.0)).reset().update(torch.tensor(2.0))
        t.assert_close(metric.sum, torch.tensor(2.0))
//  [python]
def test_reset_state_tensor(self) -> None:
        metric = DummySumMetric()
        torch.testing.assert_close(metric.sum, torch.tensor(0.0))
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        torch.testing.assert_close(metric.sum, torch.tensor(3.0))

        # reset once
        metric.reset()
        torch.testing.assert_close(metric.sum, torch.tensor(0.0))
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        torch.testing.assert_close(metric.sum, torch.tensor(3.0))

        # reset again
        metric.reset()
        torch.testing.assert_close(metric.sum, torch.tensor(0.0))

        # update and reset chained
        metric.update(torch.tensor(1.0)).reset().update(torch.tensor(2.0))
        torch.testing.assert_close(metric.sum, torch.tensor(2.0))
//  [python]
def test_reset_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric()
        t.assert_close(metric.x, [])
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        t.assert_close(metric.x, [torch.tensor(1.0), torch.tensor(2.0)])

        # reset once
        metric.reset()
        t.assert_close(metric.x, [])
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        t.assert_close(metric.x, [torch.tensor(1.0), torch.tensor(2.0)])

        # reset again
        metric.reset()
        t.assert_close(metric.x, [])

        # update and reset chained
        metric.update(torch.tensor(1.0)).reset().update(torch.tensor(2.0))
        t.assert_close(metric.x, [torch.tensor(2.0)])
//  [python]
def test_reset_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric()
        torch.testing.assert_close(metric.x, [])
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        torch.testing.assert_close(metric.x, [torch.tensor(1.0), torch.tensor(2.0)])

        # reset once
        metric.reset()
        torch.testing.assert_close(metric.x, [])
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        torch.testing.assert_close(metric.x, [torch.tensor(1.0), torch.tensor(2.0)])

        # reset again
        metric.reset()
        torch.testing.assert_close(metric.x, [])

        # update and reset chained
        metric.update(torch.tensor(1.0)).reset().update(torch.tensor(2.0))
        torch.testing.assert_close(metric.x, [torch.tensor(2.0)])
//  [python]
def test_add_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric()
        metric._add_state("x1", {})
        t.assert_close(metric.x1, {})

        metric._add_state("x2", {"doc1": torch.tensor(0.0)})
        t.assert_close(metric.x2, {"doc1": torch.tensor(0.0)})

        tensor_dict = {"doc1": torch.tensor(0.0)}
        metric._add_state("x3", tensor_dict)
        t.assert_close(metric.x3, tensor_dict)
        tensor_dict["doc1"] = torch.tensor(1.0)
        self.assertNotEqual(metric.x3, tensor_dict)
        t.assert_close(metric.x3, {"doc1": torch.tensor(0.0)})

        metric._add_state(
            "x4",
            {
                "doc1": torch.tensor(0.0),
                "doc2": torch.tensor(0.0),
                "doc3": torch.tensor(0.0),
                "doc4": torch.tensor(0.0),
            },
        )
        t.assert_close(
            metric.x4,
            {
                "doc1": torch.tensor(0.0),
                "doc2": torch.tensor(0.0),
                "doc3": torch.tensor(0.0),
                "doc4": torch.tensor(0.0),
            },
        )
//  [python]
def test_add_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric()
        metric._add_state("x1", {})
        torch.testing.assert_close(metric.x1, {})

        metric._add_state("x2", {"doc1": torch.tensor(0.0)})
        torch.testing.assert_close(metric.x2, {"doc1": torch.tensor(0.0)})

        tensor_dict = {"doc1": torch.tensor(0.0)}
        metric._add_state("x3", tensor_dict)
        torch.testing.assert_close(metric.x3, tensor_dict)
        tensor_dict["doc1"] = torch.tensor(1.0)
        self.assertNotEqual(metric.x3, tensor_dict)
        torch.testing.assert_close(metric.x3, {"doc1": torch.tensor(0.0)})

        metric._add_state(
            "x4",
            {
                "doc1": torch.tensor(0.0),
                "doc2": torch.tensor(0.0),
                "doc3": torch.tensor(0.0),
                "doc4": torch.tensor(0.0),
            },
        )
        torch.testing.assert_close(
            metric.x4,
            {
                "doc1": torch.tensor(0.0),
                "doc2": torch.tensor(0.0),
                "doc3": torch.tensor(0.0),
                "doc4": torch.tensor(0.0),
            },
        )
//  [python]
def test_reset_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric()
        t.assert_close(metric.x, {})
        metric.update("doc1", torch.tensor(0.0))
        metric.update("doc1", torch.tensor(1.0))
        t.assert_close(metric.x, {"doc1": torch.tensor(1.0)})

        # reset once
        metric.reset()
        t.assert_close(metric.x, {})
        metric.update("doc1", torch.tensor(0.0))
        metric.update("doc1", torch.tensor(1.0))
        t.assert_close(metric.x, {"doc1": torch.tensor(1.0)})

        # reset again
        metric.reset()
        t.assert_close(metric.x, {})

        # update and reset chained
        metric.update("doc1", torch.tensor(2.0)).reset().update(
            "doc1", torch.tensor(1.0)
        )
        t.assert_close(metric.x, {"doc1": torch.tensor(1.0)})

        # reset again
        metric.reset()
        t.assert_close(metric.x, {})
        metric.update("doc1", torch.tensor(0.0))
        metric.update("doc2", torch.tensor(1.0))
        t.assert_close(
            metric.x, {"doc1": torch.tensor(0.0), "doc2": torch.tensor(1.0)}
        )
//  [python]
def test_reset_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric()
        torch.testing.assert_close(metric.x, {})
        metric.update("doc1", torch.tensor(0.0))
        metric.update("doc1", torch.tensor(1.0))
        torch.testing.assert_close(metric.x, {"doc1": torch.tensor(1.0)})

        # reset once
        metric.reset()
        torch.testing.assert_close(metric.x, {})
        metric.update("doc1", torch.tensor(0.0))
        metric.update("doc1", torch.tensor(1.0))
        torch.testing.assert_close(metric.x, {"doc1": torch.tensor(1.0)})

        # reset again
        metric.reset()
        torch.testing.assert_close(metric.x, {})

        # update and reset chained
        metric.update("doc1", torch.tensor(2.0)).reset().update(
            "doc1", torch.tensor(1.0)
        )
        torch.testing.assert_close(metric.x, {"doc1": torch.tensor(1.0)})

        # reset again
        metric.reset()
        torch.testing.assert_close(metric.x, {})
        metric.update("doc1", torch.tensor(0.0))
        metric.update("doc2", torch.tensor(1.0))
        torch.testing.assert_close(
            metric.x, {"doc1": torch.tensor(0.0), "doc2": torch.tensor(1.0)}
        )
//  [python]
def test_to_device_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric().to("cuda")
        metric.update(torch.tensor(1.0).to("cuda")).compute()
        torch.testing.assert_allclose(metric.x, [torch.tensor(1.0, device="cuda")])

        metric.to("cpu")
        torch.testing.assert_allclose(metric.x, [torch.tensor(1.0, device="cpu")])
        metric.to("cuda")
        torch.testing.assert_allclose(metric.x, [torch.tensor(1.0, device="cuda")])
//  [python]
def test_to_device_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric().to("cuda")
        metric.update(torch.tensor(1.0).to("cuda")).compute()
        torch.testing.assert_close(metric.x, [torch.tensor(1.0, device="cuda")])

        metric.to("cpu")
        torch.testing.assert_close(metric.x, [torch.tensor(1.0, device="cpu")])
        metric.to("cuda")
        torch.testing.assert_close(metric.x, [torch.tensor(1.0, device="cuda")])
//  [python]
def test_save_load_state_dict_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric()
        self.assertDictEqual(metric.state_dict(), {"x": []})
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        self.assertDictEqual(
            metric.state_dict(), {"x": [torch.tensor(1.0), torch.tensor(2.0)]}
        )

        state_dict = metric.state_dict()
        loaded_metric = DummySumListStateMetric()
        loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            loaded_metric.state_dict(), {"x": [torch.tensor(1.0), torch.tensor(2.0)]}
        )
        self.assertEqual(loaded_metric.update(torch.tensor(-1.0)).compute(), torch.tensor(2.0))
        self.assertDictEqual(
            loaded_metric.state_dict(),
            {"x": [torch.tensor(1.0), torch.tensor(2.0), torch.tensor(-1.0)]},
        )

        another_loaded_metric = DummySumListStateMetric()
        another_loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            another_loaded_metric.state_dict(),
            {"x": [torch.tensor(1.0), torch.tensor(2.0)]},
        )
//  [python]
def test_save_load_state_dict_state_tensor_list(self) -> None:
        metric = DummySumListStateMetric()
        self.assertDictEqual(metric.state_dict(), {"x": []})
        metric.update(torch.tensor(1.0))
        metric.update(torch.tensor(2.0))
        self.assertDictEqual(
            metric.state_dict(), {"x": [torch.tensor(1.0), torch.tensor(2.0)]}
        )

        state_dict = metric.state_dict()
        loaded_metric = DummySumListStateMetric()
        loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            loaded_metric.state_dict(), {"x": [torch.tensor(1.0), torch.tensor(2.0)]}
        )
        torch.testing.assert_close(
            loaded_metric.update(torch.tensor(-1.0)).compute(), torch.tensor(2.0)
        )
        self.assertDictEqual(
            loaded_metric.state_dict(),
            {"x": [torch.tensor(1.0), torch.tensor(2.0), torch.tensor(-1.0)]},
        )

        another_loaded_metric = DummySumListStateMetric()
        another_loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            another_loaded_metric.state_dict(),
            {"x": [torch.tensor(1.0), torch.tensor(2.0)]},
        )
//  [python]
def test_to_device_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric().to("cuda")
        metric.update("doc1", torch.tensor(1.0).to("cuda")).compute()
        torch.testing.assert_allclose(metric.x, {"doc1": torch.tensor(1.0, device="cuda")})

        metric.to("cpu")
        torch.testing.assert_allclose(metric.x, {"doc1": torch.tensor(1.0, device="cpu")})
        metric.to("cuda")
        torch.testing.assert_allclose(metric.x, {"doc1": torch.tensor(1.0, device="cuda")})
//  [python]
def test_to_device_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric().to("cuda")
        metric.update("doc1", torch.tensor(1.0).to("cuda")).compute()
        torch.testing.assert_close(metric.x, {"doc1": torch.tensor(1.0, device="cuda")})

        metric.to("cpu")
        torch.testing.assert_close(metric.x, {"doc1": torch.tensor(1.0, device="cpu")})
        metric.to("cuda")
        torch.testing.assert_close(metric.x, {"doc1": torch.tensor(1.0, device="cuda")})
//  [python]
def test_save_load_state_dict_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric()
        self.assertDictEqual(metric.state_dict(), {"x": {}})
        metric.update("doc1", torch.tensor(1.0))
        metric.update("doc2", torch.tensor(2.0))
        self.assertDictEqual(
            metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(2.0)}},
        )

        state_dict = metric.state_dict()
        loaded_metric = DummySumDictStateMetric()
        loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            loaded_metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(2.0)}},
        )
        self.assertEqual(
            loaded_metric.update("doc2", torch.tensor(-1.0)).compute(),
            {"doc1": torch.tensor(1.0), "doc2": torch.tensor(1.0)},
        )

        self.assertDictEqual(
            loaded_metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(1.0)}},
        )

        another_loaded_metric = DummySumDictStateMetric()
        another_loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            another_loaded_metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(2.0)}},
        )
//  [python]
def test_save_load_state_dict_state_tensor_dict(self) -> None:
        metric = DummySumDictStateMetric()
        self.assertDictEqual(metric.state_dict(), {"x": {}})
        metric.update("doc1", torch.tensor(1.0))
        metric.update("doc2", torch.tensor(2.0))
        self.assertDictEqual(
            metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(2.0)}},
        )

        state_dict = metric.state_dict()
        loaded_metric = DummySumDictStateMetric()
        loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            loaded_metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(2.0)}},
        )
        torch.testing.assert_close(
            loaded_metric.update("doc2", torch.tensor(-1.0)).compute(),
            {"doc1": torch.tensor(1.0), "doc2": torch.tensor(1.0)},
        )

        self.assertDictEqual(
            loaded_metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(1.0)}},
        )

        another_loaded_metric = DummySumDictStateMetric()
        another_loaded_metric.load_state_dict(state_dict)
        self.assertDictEqual(
            another_loaded_metric.state_dict(),
            {"x": {"doc1": torch.tensor(1.0), "doc2": torch.tensor(2.0)}},
        )
//  [python]
def test_load_state_dict_strict(self) -> None:
        metric = DummySumMetric()
        metric.update(torch.tensor(3.0))

        my_state_dict = {}
        my_state_dict["nonexistent_state"] = torch.tensor(1.0)
        loaded_metric = DummySumMetric()
        with self.assertRaisesRegex(
            RuntimeError,
            r"Error\(s\) in loading state_dict for DummySumMetric. "
            r"Encountered missing keys: {'sum'} and unexpected keys: "
            r"{'nonexistent_state'}.",
        ):
            loaded_metric.load_state_dict(my_state_dict)

        loaded_metric.load_state_dict(my_state_dict, strict=False)
        t.assert_close(loaded_metric.sum, torch.tensor(0.0))

        state_dict = metric.state_dict()
        state_dict["nonexistent_state"] = torch.tensor(1.0)
        with self.assertRaisesRegex(
            RuntimeError,
            r"Error\(s\) in loading state_dict for DummySumMetric. "
            r"Encountered missing keys: set\(\) and unexpected keys: "
            r"{'nonexistent_state'}.",
        ):
            loaded_metric.load_state_dict(state_dict)

        loaded_metric.load_state_dict(state_dict, strict=False)
        t.assert_close(loaded_metric.sum, torch.tensor(3.0))
//  [python]
def test_load_state_dict_strict(self) -> None:
        metric = DummySumMetric()
        metric.update(torch.tensor(3.0))

        my_state_dict = {}
        my_state_dict["nonexistent_state"] = torch.tensor(1.0)
        loaded_metric = DummySumMetric()
        with self.assertRaisesRegex(
            RuntimeError,
            r"Error\(s\) in loading state_dict for DummySumMetric. "
            r"Encountered missing keys: {'sum'} and unexpected keys: "
            r"{'nonexistent_state'}.",
        ):
            loaded_metric.load_state_dict(my_state_dict)

        loaded_metric.load_state_dict(my_state_dict, strict=False)
        torch.testing.assert_close(loaded_metric.sum, torch.tensor(0.0))

        state_dict = metric.state_dict()
        state_dict["nonexistent_state"] = torch.tensor(1.0)
        with self.assertRaisesRegex(
            RuntimeError,
            r"Error\(s\) in loading state_dict for DummySumMetric. "
            r"Encountered missing keys: set\(\) and unexpected keys: "
            r"{'nonexistent_state'}.",
        ):
            loaded_metric.load_state_dict(state_dict)

        loaded_metric.load_state_dict(state_dict, strict=False)
        torch.testing.assert_close(loaded_metric.sum, torch.tensor(3.0))
//  [python]
def test_ws_calculation():
    """Check if convolution weights are properly standardized."""
    model = SimpleConvModel()
    apply_weight_standardization(module=model)
    var1, mean1 = torch.var_mean(model.conv1.weight, dim=[1, 2, 3], unbiased=False)
    var2, mean2 = torch.var_mean(model.conv2.weight, dim=[1, 2, 3], unbiased=False)
    torch.testing.assert_allclose(var1, torch.ones_like(var1))
    torch.testing.assert_allclose(var2, torch.ones_like(var2))
    torch.testing.assert_allclose(mean1, torch.zeros_like(mean1))
    torch.testing.assert_allclose(mean2, torch.zeros_like(mean2))
//  [python]
def test_ws_calculation():
    """Check if convolution weights are properly standardized."""
    model = SimpleConvModel()
    apply_weight_standardization(module=model)
    var1, mean1 = torch.var_mean(model.conv1.weight, dim=[1, 2, 3], unbiased=False)
    var2, mean2 = torch.var_mean(model.conv2.weight, dim=[1, 2, 3], unbiased=False)
    torch.testing.assert_close(var1, torch.ones_like(var1))
    torch.testing.assert_close(var2, torch.ones_like(var2))
    torch.testing.assert_close(mean1, torch.zeros_like(mean1))
    torch.testing.assert_close(mean2, torch.zeros_like(mean2))