jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testScaledDotProductAttention():
    from deepchem.models import ScaledDotProductAttention as SDPA
    attn = SDPA()
    x = torch.ones(2, 5)
    # Linear layers for making query, key, value
    Q, K, V = nn.Parameter(torch.ones(5)), nn.Parameter(
        torch.ones(5)), nn.Parameter(torch.ones(5))
    query, key, value = Q * x, K * x, V * x
    mask = torch.Tensor([1, 0])
    x_out, attn_score = attn(query, key, value, mask=mask)
    torch.testing.assert_close(x_out, torch.ones(2, 5))
    torch.testing.assert_close(attn_score, torch.Tensor([[1, 0], [1, 0]]))
//  [python]
def test_multidimensional_support_GaussianKernel(self):
        # Gaussian 2D data.
        train_Xs = torch.normal(torch.zeros((100, 2)), torch.ones((100, 2)))
        kernel = models.GaussianKernel(bandwidth=1.0)
        model = models.KernelDensityEstimator(train_Xs, kernel)

        # Build meshgrid.
        dx = 0.1
        X = torch.arange(-8, 8, dx)
        Y = torch.arange(-8, 8, dx)
        xx, yy = torch.meshgrid(X, Y, indexing="ij")
        meshgrid = torch.stack((xx, yy), axis=2).view(-1, 2)

        log_probs = model(meshgrid)
        integral = torch.sum(torch.exp(log_probs) * dx * dx)
        torch.testing.assert_allclose(integral, torch.tensor(1.0))
//  [python]
def test_multidimensional_support_GaussianKernel(self):
        # Gaussian 2D data.
        train_Xs = torch.normal(torch.zeros((100, 2)), torch.ones((100, 2)))
        kernel = models.GaussianKernel(bandwidth=1.0)
        model = models.KernelDensityEstimator(train_Xs, kernel)

        # Build meshgrid.
        dx = 0.1
        X = torch.arange(-8, 8, dx)
        Y = torch.arange(-8, 8, dx)
        xx, yy = torch.meshgrid(X, Y, indexing="ij")
        meshgrid = torch.stack((xx, yy), axis=2).view(-1, 2)

        log_probs = model(meshgrid)
        integral = torch.sum(torch.exp(log_probs) * dx * dx)
        torch.testing.assert_close(integral, torch.tensor(1.0))
//  [python]
def test_multidimensional_support_ParzenWindowKernel(self):
        # Gaussian 2D data.
        train_Xs = torch.normal(torch.zeros((100, 2)), torch.ones((100, 2)))
        kernel = models.ParzenWindowKernel(bandwidth=1.0)
        model = models.KernelDensityEstimator(train_Xs, kernel)

        # Build meshgrid.
        dx = 0.1
        X = torch.arange(-8, 8, dx)
        Y = torch.arange(-8, 8, dx)
        xx, yy = torch.meshgrid(X, Y, indexing="ij")
        meshgrid = torch.stack((xx, yy), axis=2).view(-1, 2)

        log_probs = model(meshgrid)
        integral = torch.sum(torch.exp(log_probs) * dx * dx)
        self.assertTrue(torch.allclose(integral, torch.tensor(1.0)))
//  [python]
def test_multidimensional_support_ParzenWindowKernel(self):
        # Gaussian 2D data.
        train_Xs = torch.normal(torch.zeros((100, 2)), torch.ones((100, 2)))
        kernel = models.ParzenWindowKernel(bandwidth=1.0)
        model = models.KernelDensityEstimator(train_Xs, kernel)

        # Build meshgrid.
        dx = 0.1
        X = torch.arange(-8, 8, dx)
        Y = torch.arange(-8, 8, dx)
        xx, yy = torch.meshgrid(X, Y, indexing="ij")
        meshgrid = torch.stack((xx, yy), axis=2).view(-1, 2)

        log_probs = model(meshgrid)
        integral = torch.sum(torch.exp(log_probs) * dx * dx)
        torch.testing.assert_close(integral, torch.tensor(1.0))
//  [python]
def test_save_load_safetensors_used(self, net_fit, file_str, tmpdir):
        # Safetensors' capacity to save and load net params is already covered
        # in other tests. This is a test to exclude the (trivial) bug that even
        # with use_safetensors=True, safetensors is not actually being used
        # (instead accidentally using pickle or something like that). To test
        # this, we directly open the stored file using safetensors and check its
        # contents. If it were, say, a pickle file, this test would fail.
        from safetensors import safe_open

        p = tmpdir.mkdir('skorch').join('testmodel.safetensors')

        if file_str:
            net_fit.save_params(f_params=str(p), use_safetensors=True)
        else:
            with open(str(p), 'wb') as f:
                net_fit.save_params(f_params=f, use_safetensors=True)

        state_dict_loaded = {}
        with safe_open(str(p), framework='pt', device=net_fit.device) as f:
            for key in f.keys():
                state_dict_loaded[key] = f.get_tensor(key)

        state_dict = net_fit.module_.state_dict()
        assert state_dict_loaded.keys() == state_dict.keys()
        for key in state_dict:
            t.assert_close(state_dict[key], state_dict_loaded[key])
//  [python]
def test_save_load_safetensors_used(self, net_fit, file_str, tmpdir):
        # Safetensors' capacity to save and load net params is already covered
        # in other tests. This is a test to exclude the (trivial) bug that even
        # with use_safetensors=True, safetensors is not actually being used
        # (instead accidentally using pickle or something like that). To test
        # this, we directly open the stored file using safetensors and check its
        # contents. If it were, say, a pickle file, this test would fail.
        from safetensors import safe_open

        p = tmpdir.mkdir('skorch').join('testmodel.safetensors')

        if file_str:
            net_fit.save_params(f_params=str(p), use_safetensors=True)
        else:
            with open(str(p), 'wb') as f:
                net_fit.save_params(f_params=f, use_safetensors=True)

        state_dict_loaded = {}
        with safe_open(str(p), framework='pt', device=net_fit.device) as f:
            for key in f.keys():
                state_dict_loaded[key] = f.get_tensor(key)

        state_dict = net_fit.module_.state_dict()
        assert state_dict_loaded.keys() == state_dict.keys()
        for key in state_dict:
            torch.testing.assert_close(state_dict[key], state_dict_loaded[key])
//  [python]
def test_rematerialization(executor, device, _):
    n_fusion_regions = 7

    @value_and_grad
    def func(t0):
        for _ in range(n_fusion_regions):
            t1 = ttorch.cos(t0)
            t2 = ttorch.sin(t1)
            t3 = ttorch.cos(t2)
            t4 = ttorch.matmul(t3, t3)  # Fusion breaks here
            t0 = t4
        return t4

    t0 = make_tensor(2, 2, dtype=torch.float32, device=device)

    # Result with rematerialization and without rematerialization should match
    result_with_remat = thunder.compile(
        func,
        disable_preprocessing=True,
        use_rematerialization=True,
    )(t0)
    assert not isinstance(result_with_remat, Exception)

    result_without_remat = thunder.compile(
        func,
        disable_preprocessing=True,
        use_rematerialization=False,
    )(t0)

    torch.testing.assert_allclose(result_with_remat, result_without_remat)
//  [python]
def test_rematerialization(executor, device, _):
    n_fusion_regions = 7

    @value_and_grad
    def func(t0):
        for _ in range(n_fusion_regions):
            t1 = ttorch.cos(t0)
            t2 = ttorch.sin(t1)
            t3 = ttorch.cos(t2)
            t4 = ttorch.matmul(t3, t3)  # Fusion breaks here
            t0 = t4
        return t4

    t0 = make_tensor(2, 2, dtype=torch.float32, device=device)

    # Result with rematerialization and without rematerialization should match
    result_with_remat = thunder.compile(
        func,
        disable_preprocessing=True,
        use_rematerialization=True,
    )(t0)
    assert not isinstance(result_with_remat, Exception)

    result_without_remat = thunder.compile(
        func,
        disable_preprocessing=True,
        use_rematerialization=False,
    )(t0)

    torch.testing.assert_close(result_with_remat, result_without_remat)
//  [python]
def test_masks(masks):
    rles_0 = mask_to_rle_pytorch(masks)
    rles_2 = mask_to_rle_pytorch_2(masks)
    
    for i in range(len(rles_0)):
        torch.testing.assert_allclose(torch.tensor(rles_0[i]['counts']), torch.tensor(rles_2[i]['counts']))
//  [python]
def test_masks(masks):
    rles_0 = mask_to_rle_pytorch(masks)
    rles_2 = mask_to_rle_pytorch_2(masks)
    
    for i in range(len(rles_0)):
        torch.testing.assert_close(torch.tensor(rles_0[i]['counts']), torch.tensor(rles_2[i]['counts']))
//  [python]
def assert_close(a: Tensor, b: Union[Tensor, List[Any]]) -> None:
    """Assert that ``a`` and ``b`` are element-wise equal within a tolerance."""
    if not isinstance(b, Tensor):
        b = torch.tensor(b, device=device, dtype=a.dtype)

    torch.testing.assert_allclose(a, b)
//  [python]
def assert_close(a: Tensor, b: Union[Tensor, List[Any]]) -> None:
    """Assert that ``a`` and ``b`` are element-wise equal within a tolerance."""
    if not isinstance(b, Tensor):
        b = torch.tensor(b, device=device, dtype=a.dtype)

    torch.testing.assert_close(a, b)
//  [python]
def test_lora_wrappers_llama_works() -> None:
    # Construct a smaller LLaMAModel to prevent CI from failing
    llama_config = LLaMAConfig(
        model_dim=1024,
        max_seq_len=2048,
        vocab_info=VocabularyInfo(
            size=32000, unk_idx=0, bos_idx=1, eos_idx=2, pad_idx=None
        ),
        num_layers=16,
        num_attn_heads=8,
        num_key_value_heads=8,
        ffn_inner_dim=1024 * 4,
        ffn_inner_dim_to_multiple=1,
        dropout_p=0.1,
    )
    model = create_llama_model(llama_config, device=torch.device("cpu"))

    lora_config = get_llama_lora_config()

    inputs = torch.LongTensor([[1, 2], [1, 3]], device=torch.device("cpu"))

    model.eval()

    with torch.inference_mode():
        output_before_wrap, _ = model.decode(seqs=inputs, padding_mask=None)

    model = wrap_lora(model, lora_config)  # type: ignore[assignment]

    with torch.inference_mode():
        output_after_wrap, _ = model.decode(seqs=inputs, padding_mask=None)

    # Outputs should be the same as lora_B is initialized with zeros.
    t.assert_close(output_before_wrap, output_after_wrap)

    model = unwrap_lora(model, merge=False)  # type: ignore[assignment]

    with torch.inference_mode():
        output_after_unwrap, _ = model.decode(seqs=inputs, padding_mask=None)

    t.assert_close(output_after_wrap, output_after_unwrap)

    model = wrap_lora(model, lora_config)  # type: ignore[assignment]
    merge_lora(model)

    with torch.inference_mode():
        output_after_merge, _ = model.decode(seqs=inputs, padding_mask=None)

    t.assert_close(output_after_unwrap, output_after_merge)

    unmerge_lora(model)

    with torch.inference_mode():
        output_after_unmerge, _ = model.decode(seqs=inputs, padding_mask=None)

    t.assert_close(output_after_merge, output_after_unmerge)

    model.train()
    freeze_non_lora(model, unfreeze_bias="none")

    for name, param in model.named_parameters():
        if param.requires_grad:
            assert "lora_" in name
//  [python]
def test_lora_wrappers_llama_works() -> None:
    # Construct a smaller LLaMAModel to prevent CI from failing
    llama_config = LLaMAConfig(
        model_dim=1024,
        max_seq_len=2048,
        vocab_info=VocabularyInfo(
            size=32000, unk_idx=0, bos_idx=1, eos_idx=2, pad_idx=None
        ),
        num_layers=16,
        num_attn_heads=8,
        num_key_value_heads=8,
        ffn_inner_dim=1024 * 4,
        ffn_inner_dim_to_multiple=1,
        dropout_p=0.1,
    )
    model = create_llama_model(llama_config, device=torch.device("cpu"))

    lora_config = get_llama_lora_config()

    inputs = torch.LongTensor([[1, 2], [1, 3]], device=torch.device("cpu"))

    model.eval()

    with torch.inference_mode():
        output_before_wrap, _ = model.decode(seqs=inputs, padding_mask=None)

    model = wrap_lora(model, lora_config)  # type: ignore[assignment]

    with torch.inference_mode():
        output_after_wrap, _ = model.decode(seqs=inputs, padding_mask=None)

    # Outputs should be the same as lora_B is initialized with zeros.
    torch.testing.assert_close(output_before_wrap, output_after_wrap)

    model = unwrap_lora(model, merge=False)  # type: ignore[assignment]

    with torch.inference_mode():
        output_after_unwrap, _ = model.decode(seqs=inputs, padding_mask=None)

    torch.testing.assert_close(output_after_wrap, output_after_unwrap)

    model = wrap_lora(model, lora_config)  # type: ignore[assignment]
    merge_lora(model)

    with torch.inference_mode():
        output_after_merge, _ = model.decode(seqs=inputs, padding_mask=None)

    torch.testing.assert_close(output_after_unwrap, output_after_merge)

    unmerge_lora(model)

    with torch.inference_mode():
        output_after_unmerge, _ = model.decode(seqs=inputs, padding_mask=None)

    torch.testing.assert_close(output_after_merge, output_after_unmerge)

    model.train()
    freeze_non_lora(model, unfreeze_bias="none")

    for name, param in model.named_parameters():
        if param.requires_grad:
            assert "lora_" in name
//  [python]
def test_retrieval_precision_single_updates_1_queries_avg_none_other_example(
        self,
    ) -> None:
        metric = RetrievalPrecision(k=3)
        metric.update(
            torch.tensor([0.1, 0.2, 0.3, 0.3, 0.4]), torch.tensor([1, 0, 0, 1, 1])
        )
        metric.update(
            torch.tensor([0.9, 0.9, 0.3, 0.1, 0.1]), torch.tensor([1, 1, 0, 0, 0])
        )

        expected_result = torch.tensor([1.0])
        actual_result = metric.compute()
        torch.testing.assert_allclose(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_1_queries_avg_none_other_example(
        self,
    ) -> None:
        metric = RetrievalPrecision(k=3)
        metric.update(
            torch.tensor([0.1, 0.2, 0.3, 0.3, 0.4]), torch.tensor([1, 0, 0, 1, 1])
        )
        metric.update(
            torch.tensor([0.9, 0.9, 0.3, 0.1, 0.1]), torch.tensor([1, 1, 0, 0, 0])
        )

        expected_result = torch.tensor([1.0])
        actual_result = metric.compute()
        torch.testing.assert_close(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_1_query_avg_none(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        metric = RetrievalPrecision(k=2)
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([0.5])
        torch.testing.assert_allclose(actual_result, expected_result)

        # reset metric and compute a new value with a new target
        metric.reset()

        target[2] = 0
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([0.0])
        torch.testing.assert_allclose(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_1_query_avg_none(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        metric = RetrievalPrecision(k=2)
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([0.5])
        torch.testing.assert_close(actual_result, expected_result)

        # reset metric and compute a new value with a new target
        metric.reset()

        target[2] = 0
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([0.0])
        torch.testing.assert_close(actual_result, expected_result)
//  [python]
def run_single_experiment(config: ExperimentConfig) -> ExperimentResults:
    device = torch.device("cuda")
    composite_mha = CompositeMHA(
        config.num_heads, config.embed_dim, device, config.dtype
    )
    composite_mha.reset_parameters()
    query, key, value = generate_inputs(
        config.batch_size,
        config.q_seq_len,
        config.k_seq_len,
        config.embed_dim,
        config.dtype,
        device,
    )
    attn_mask = CausalBias(
        CausalVariant.LOWER_RIGHT, config.q_seq_len, config.k_seq_len
    )
    attn_mask_tensor = attn_mask._materialize(device)

    materialized_mask_time = benchmark_torch_function_in_microseconds(
        composite_mha, query, key, value, attn_mask_tensor
    )
    attn_mask_subclass_time = benchmark_torch_function_in_microseconds(
        composite_mha, query, key, value, attn_mask
    )
    torch.testing.assert_allclose(
        composite_mha(query, key, value, attn_mask_tensor),
        composite_mha(query, key, value, attn_mask),
    )

    return ExperimentResults(
        materialized_mask_time=materialized_mask_time,
        attn_mask_subclass_time=attn_mask_subclass_time,
    )
//  [python]
def run_single_experiment(config: ExperimentConfig) -> ExperimentResults:
    device = torch.device("cuda")
    composite_mha = CompositeMHA(
        config.num_heads, config.embed_dim, device, config.dtype
    )
    composite_mha.reset_parameters()
    query, key, value = generate_inputs(
        config.batch_size,
        config.q_seq_len,
        config.k_seq_len,
        config.embed_dim,
        config.dtype,
        device,
    )
    attn_mask = CausalBias(
        CausalVariant.LOWER_RIGHT, config.q_seq_len, config.k_seq_len
    )
    attn_mask_tensor = attn_mask._materialize(device)

    materialized_mask_time = benchmark_torch_function_in_microseconds(
        composite_mha, query, key, value, attn_mask_tensor
    )
    attn_mask_subclass_time = benchmark_torch_function_in_microseconds(
        composite_mha, query, key, value, attn_mask
    )
    torch.testing.assert_close(
        composite_mha(query, key, value, attn_mask_tensor),
        composite_mha(query, key, value, attn_mask),
    )

    return ExperimentResults(
        materialized_mask_time=materialized_mask_time,
        attn_mask_subclass_time=attn_mask_subclass_time,
    )
//  [python]
def test_retrieval_precision_single_updates_1_query_avg_macro(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        metric = RetrievalPrecision(k=2, avg="macro")
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor(0.5)
        t.assert_close(actual_result, expected_result)

        # reset metric and compute a new value with a new target
        metric.reset()

        target[2] = 0
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor(0.0)
        t.assert_close(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_1_query_avg_macro(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        metric = RetrievalPrecision(k=2, avg="macro")
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor(0.5)
        torch.testing.assert_close(actual_result, expected_result)

        # reset metric and compute a new value with a new target
        metric.reset()

        target[2] = 0
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor(0.0)
        torch.testing.assert_close(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_n_queries_avg_none(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        metric = RetrievalPrecision(k=2, num_queries=4)
        indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
        metric.update(input, target, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1 / 2, 1 / 2, torch.nan, torch.nan])
        torch.testing.assert_allclose(actual_result, expected_result, equal_nan=True)

        input2 = torch.tensor([0.4, 0.1, 0.6, 0.8, 0.7, 0.9, 0.3])
        target2 = torch.tensor([1, 0, 1, 0, 1, 1, 0])
        indexes = torch.tensor([2, 2, 2, 3, 3, 3, 3])
        metric.update(input2, target2, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1 / 2, 1 / 2, 1.0, 1 / 2])
        torch.testing.assert_allclose(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_n_queries_avg_none(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        metric = RetrievalPrecision(k=2, num_queries=4)
        indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
        metric.update(input, target, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1 / 2, 1 / 2, torch.nan, torch.nan])
        torch.testing.assert_close(actual_result, expected_result, equal_nan=True)

        input2 = torch.tensor([0.4, 0.1, 0.6, 0.8, 0.7, 0.9, 0.3])
        target2 = torch.tensor([1, 0, 1, 0, 1, 1, 0])
        indexes = torch.tensor([2, 2, 2, 3, 3, 3, 3])
        metric.update(input2, target2, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1 / 2, 1 / 2, 1.0, 1 / 2])
        torch.testing.assert_close(actual_result, expected_result)
//  [python]
def test_retrieval_precision_empty_target_action(self) -> None:
        input = torch.tensor([0.2, 0.3])
        target = torch.tensor([0, 0])

        metric = RetrievalPrecision(k=2, empty_target_action="neg")
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([0.0])
        t.assert_close(actual_result, expected_result)

        metric = RetrievalPrecision(k=2, empty_target_action="pos")
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([1.0])
        t.assert_close(actual_result, expected_result)

        metric = RetrievalPrecision(k=2, empty_target_action="err")
        with self.assertRaisesRegex(
            ValueError, r"no positive value found in target=tensor\(\[0\.\, 0\.\]\)\."
        ):
            metric.update(input, target)
            metric.compute()
//  [python]
def test_retrieval_precision_empty_target_action(self) -> None:
        input = torch.tensor([0.2, 0.3])
        target = torch.tensor([0, 0])

        metric = RetrievalPrecision(k=2, empty_target_action="neg")
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([0.0])
        torch.testing.assert_close(actual_result, expected_result)

        metric = RetrievalPrecision(k=2, empty_target_action="pos")
        metric.update(input, target)
        actual_result = metric.compute()
        expected_result = torch.tensor([1.0])
        torch.testing.assert_close(actual_result, expected_result)

        metric = RetrievalPrecision(k=2, empty_target_action="err")
        with self.assertRaisesRegex(
            ValueError, r"no positive value found in target=tensor\(\[0\.\, 0\.\]\)\."
        ):
            metric.update(input, target)
            metric.compute()
//  [python]
def test_retrieval_precision_single_updates_n_queries_avg_macro(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])
        indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])

        metric = RetrievalPrecision(k=2, num_queries=4, avg="macro")
        metric.update(input, target, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor(0.5)
        t.assert_close(actual_result, expected_result)

        input2 = torch.tensor([0.4, 0.1, 0.6, 0.8, 0.7, 0.9, 0.3])
        target2 = torch.tensor([1, 0, 1, 0, 1, 1, 0])
        indexes = torch.tensor([2, 2, 2, 3, 3, 3, 3])
        # add value (1 + 1 / 2) / 2 = 0.750 to scores
        metric.update(input2, target2, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor((0.5 + 0.750) / 2)
        t.assert_close(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_n_queries_avg_macro(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])
        indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])

        metric = RetrievalPrecision(k=2, num_queries=4, avg="macro")
        metric.update(input, target, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor(0.5)
        torch.testing.assert_close(actual_result, expected_result)

        input2 = torch.tensor([0.4, 0.1, 0.6, 0.8, 0.7, 0.9, 0.3])
        target2 = torch.tensor([1, 0, 1, 0, 1, 1, 0])
        indexes = torch.tensor([2, 2, 2, 3, 3, 3, 3])
        # add value (1 + 1 / 2) / 2 = 0.750 to scores
        metric.update(input2, target2, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor((0.5 + 0.750) / 2)
        torch.testing.assert_close(actual_result, expected_result)
//  [python]
def test_retrieval_precision_no_update(self) -> None:
        metric = RetrievalPrecision()
        expected_result = torch.tensor([torch.nan])
        actual_result = metric.compute()
        torch.testing.assert_allclose(actual_result, expected_result, equal_nan=True)

        metric = RetrievalPrecision(num_queries=2)
        expected_result = torch.tensor([torch.nan] * 2)
        actual_result = metric.compute()
        torch.testing.assert_allclose(actual_result, expected_result, equal_nan=True)
//  [python]
def test_retrieval_precision_no_update(self) -> None:
        metric = RetrievalPrecision()
        expected_result = torch.tensor([torch.nan])
        actual_result = metric.compute()
        torch.testing.assert_close(actual_result, expected_result, equal_nan=True)

        metric = RetrievalPrecision(num_queries=2)
        expected_result = torch.tensor([torch.nan] * 2)
        actual_result = metric.compute()
        torch.testing.assert_close(actual_result, expected_result, equal_nan=True)
//  [python]
def test_noisy(layer_class, device, seed=0):
    torch.manual_seed(seed)
    layer = layer_class(3, 4, device=device)
    x = torch.randn(10, 3, device=device)
    y1 = layer(x)
    layer.reset_noise()
    y2 = layer(x)
    y3 = layer(x)
    torch.testing.assert_allclose(y2, y3)
    with pytest.raises(AssertionError):
        torch.testing.assert_allclose(y1, y2)
//  [python]
def test_noisy(layer_class, device, seed=0):
    torch.manual_seed(seed)
    layer = layer_class(3, 4, device=device)
    x = torch.randn(10, 3, device=device)
    y1 = layer(x)
    layer.reset_noise()
    y2 = layer(x)
    y3 = layer(x)
    torch.testing.assert_close(y2, y3)
    with pytest.raises(AssertionError):
        torch.testing.assert_close(y1, y2)
//  [python]
def test_retrieval_precision_single_updates_n_queries_avg_none_docstring_example(
        self,
    ) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])
        indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
        metric = RetrievalPrecision(k=2, num_queries=2)

        metric.update(input, target, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1 / 2, 1 / 2])
        torch.testing.assert_allclose(actual_result, expected_result)

        target2 = torch.tensor([1, 0, 1, 0, 1, 1, 0])
        input2 = torch.tensor([0.4, 0.1, 0.6, 0.8, 0.7, 0.9, 0.3])
        metric.update(input2, target2, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1.0, 1 / 2])
        torch.testing.assert_allclose(actual_result, expected_result)
//  [python]
def test_retrieval_precision_single_updates_n_queries_avg_none_docstring_example(
        self,
    ) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])
        indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
        metric = RetrievalPrecision(k=2, num_queries=2)

        metric.update(input, target, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1 / 2, 1 / 2])
        torch.testing.assert_close(actual_result, expected_result)

        target2 = torch.tensor([1, 0, 1, 0, 1, 1, 0])
        input2 = torch.tensor([0.4, 0.1, 0.6, 0.8, 0.7, 0.9, 0.3])
        metric.update(input2, target2, indexes)
        actual_result = metric.compute()
        expected_result = torch.tensor([1.0, 1 / 2])
        torch.testing.assert_close(actual_result, expected_result)
//  [python]
def test_survival_function():
    random = numpy.random.default_rng(0xDEADBEEF)

    batch_size = 32

    loc = torch.tensor(numpy.array([random.standard_normal(batch_size)]))

    scale = torch.tensor(numpy.array([random.random(batch_size) + 1.0]))

    x = torch.linspace(-8.0, 8.0, batch_size).to(torch.float64)

    torch.testing.assert_allclose(
        survival_function(x, Normal(loc, scale)),
        torch.tensor(scipy.stats.norm(loc, scale).sf(x), dtype=torch.float32),
    )
//  [python]
def test_survival_function():
    random = numpy.random.default_rng(0xDEADBEEF)

    batch_size = 32

    loc = torch.tensor(numpy.array([random.standard_normal(batch_size)]))

    scale = torch.tensor(numpy.array([random.random(batch_size) + 1.0]))

    x = torch.linspace(-8.0, 8.0, batch_size).to(torch.float64)

    torch.testing.assert_close(
        survival_function(x, Normal(loc, scale)),
        torch.tensor(scipy.stats.norm(loc, scale).sf(x), dtype=torch.float32),
    )
//  [python]
def test_tanh_atanh(use_vmap, scale):
    if use_vmap:
        try:
            from torch import vmap
        except ImportError:
            try:
                from functorch import vmap
            except ImportError:
                raise pytest.skip("functorch not found")

    torch.manual_seed(0)
    x = (torch.randn(10, dtype=torch.double) * scale).requires_grad_(True)
    if not use_vmap:
        y = safetanh(x, 1e-6)
    else:
        y = vmap(safetanh, (0, None))(x, 1e-6)

    if not use_vmap:
        xp = safeatanh(y, 1e-6)
    else:
        xp = vmap(safeatanh, (0, None))(y, 1e-6)

    xp.sum().backward()
    torch.testing.assert_allclose(x.grad, torch.ones_like(x))
//  [python]
def test_tanh_atanh(use_vmap, scale):
    if use_vmap:
        try:
            from torch import vmap
        except ImportError:
            try:
                from functorch import vmap
            except ImportError:
                raise pytest.skip("functorch not found")

    torch.manual_seed(0)
    x = (torch.randn(10, dtype=torch.double) * scale).requires_grad_(True)
    if not use_vmap:
        y = safetanh(x, 1e-6)
    else:
        y = vmap(safetanh, (0, None))(x, 1e-6)

    if not use_vmap:
        xp = safeatanh(y, 1e-6)
    else:
        xp = vmap(safeatanh, (0, None))(y, 1e-6)

    xp.sum().backward()
    torch.testing.assert_close(x.grad, torch.ones_like(x))
//  [python]
def test_lstm_net_nobatch(device, out_features, hidden_size):
    time_steps = 6
    in_features = 4
    net = LSTMNet(
        out_features,
        {"input_size": hidden_size, "hidden_size": hidden_size},
        {"out_features": hidden_size},
        device=device,
    )
    # test single step vs multi-step
    x = torch.randn(time_steps, in_features, device=device)
    x_unbind = x.unbind(0)
    tds_loop = []
    hidden0_in, hidden1_in, hidden0_out, hidden1_out = [
        None,
    ] * 4
    for _x in x_unbind:
        y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(
            _x, hidden0_out, hidden1_out
        )
        td = TensorDict(
            {
                "y": y,
                "hidden0_in": hidden0_in,
                "hidden1_in": hidden1_in,
                "hidden0_out": hidden0_out,
                "hidden1_out": hidden1_out,
            },
            [],
        )
        tds_loop.append(td)
    tds_loop = torch.stack(tds_loop, 0)

    y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(x.unsqueeze(0))
    tds_vec = TensorDict(
        {
            "y": y,
            "hidden0_in": hidden0_in,
            "hidden1_in": hidden1_in,
            "hidden0_out": hidden0_out,
            "hidden1_out": hidden1_out,
        },
        [1, time_steps],
    ).squeeze(0)
    torch.testing.assert_allclose(tds_vec["y"], tds_loop["y"])
    torch.testing.assert_allclose(tds_vec["hidden0_out"][-1], tds_loop["hidden0_out"][-1])
    torch.testing.assert_allclose(tds_vec["hidden1_out"][-1], tds_loop["hidden1_out"][-1])
//  [python]
def test_lstm_net_nobatch(device, out_features, hidden_size):
    time_steps = 6
    in_features = 4
    net = LSTMNet(
        out_features,
        {"input_size": hidden_size, "hidden_size": hidden_size},
        {"out_features": hidden_size},
        device=device,
    )
    # test single step vs multi-step
    x = torch.randn(time_steps, in_features, device=device)
    x_unbind = x.unbind(0)
    tds_loop = []
    hidden0_in, hidden1_in, hidden0_out, hidden1_out = [
        None,
    ] * 4
    for _x in x_unbind:
        y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(
            _x, hidden0_out, hidden1_out
        )
        td = TensorDict(
            {
                "y": y,
                "hidden0_in": hidden0_in,
                "hidden1_in": hidden1_in,
                "hidden0_out": hidden0_out,
                "hidden1_out": hidden1_out,
            },
            [],
        )
        tds_loop.append(td)
    tds_loop = torch.stack(tds_loop, 0)

    y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(x.unsqueeze(0))
    tds_vec = TensorDict(
        {
            "y": y,
            "hidden0_in": hidden0_in,
            "hidden1_in": hidden1_in,
            "hidden0_out": hidden0_out,
            "hidden1_out": hidden1_out,
        },
        [1, time_steps],
    ).squeeze(0)
    torch.testing.assert_close(tds_vec["y"], tds_loop["y"])
    torch.testing.assert_close(tds_vec["hidden0_out"][-1], tds_loop["hidden0_out"][-1])
    torch.testing.assert_close(tds_vec["hidden1_out"][-1], tds_loop["hidden1_out"][-1])
//  [python]
def assert_close(a: torch.Tensor, b: torch.Tensor):
    """Given two Tensors, compare with a reasonable precision.

    If the dtypes mismatch, use a custom rule to cast one or the other
    """
    # increasing order of precision
    precedence = [torch.bfloat16, torch.float16, torch.float32, torch.float64]

    # demote inputs so we use the more permissive test
    if precedence.index(a.dtype) < precedence.index(b.dtype):
        b = b.type(a.dtype)
    else:
        a = a.type(b.dtype)

    if a.dtype in [torch.float16, torch.bfloat16]:
        # torch.nn.InstanceNorm3d fails rtol=6, atols=4e-2 for half precision
        torch.testing.assert_allclose(a, b, rtol=10, atol=5e-2)
    else:  # use default tolerance
        torch.testing.assert_allclose(a, b)
//  [python]
def assert_close(a: torch.Tensor, b: torch.Tensor):
    """Given two Tensors, compare with a reasonable precision.

    If the dtypes mismatch, use a custom rule to cast one or the other
    """
    # increasing order of precision
    precedence = [torch.bfloat16, torch.float16, torch.float32, torch.float64]

    # demote inputs so we use the more permissive test
    if precedence.index(a.dtype) < precedence.index(b.dtype):
        b = b.type(a.dtype)
    else:
        a = a.type(b.dtype)

    if a.dtype in [torch.float16, torch.bfloat16]:
        # torch.nn.InstanceNorm3d fails rtol=6, atols=4e-2 for half precision
        torch.testing.assert_close(a, b, rtol=10, atol=5e-2)
    else:  # use default tolerance
        torch.testing.assert_close(a, b)
//  [python]
def test_python_gru_cell(device, bias):
    gru_cell1 = GRUCell(10, 20, device=device, bias=bias)
    gru_cell2 = nn.GRUCell(10, 20, device=device, bias=bias)

    gru_cell2.load_state_dict(gru_cell1.state_dict())

    # Make sure parameters match
    for (k1, v1), (k2, v2) in zip(
        gru_cell1.named_parameters(), gru_cell2.named_parameters()
    ):
        assert k1 == k2, f"Parameter names do not match: {k1} != {k2}"
        assert (v1 == v2).all()
        assert (
            v1.shape == v2.shape
        ), f"Parameter shapes do not match: {k1} shape {v1.shape} != {k2} shape {v2.shape}"

    # Run loop
    input = torch.randn(2, 3, 10, device=device)
    h0 = torch.zeros(3, 20, device=device)
    with torch.no_grad():
        for i in range(input.size()[0]):
            h1 = gru_cell1(input[i], h0)
            h2 = gru_cell2(input[i], h0)

            # Make sure the final hidden states have the same shape
            assert h1.shape == h2.shape
            torch.testing.assert_allclose(h1, h2)
            h0 = h1
//  [python]
def test_python_gru_cell(device, bias):
    gru_cell1 = GRUCell(10, 20, device=device, bias=bias)
    gru_cell2 = nn.GRUCell(10, 20, device=device, bias=bias)

    gru_cell2.load_state_dict(gru_cell1.state_dict())

    # Make sure parameters match
    for (k1, v1), (k2, v2) in zip(
        gru_cell1.named_parameters(), gru_cell2.named_parameters()
    ):
        assert k1 == k2, f"Parameter names do not match: {k1} != {k2}"
        assert (v1 == v2).all()
        assert (
            v1.shape == v2.shape
        ), f"Parameter shapes do not match: {k1} shape {v1.shape} != {k2} shape {v2.shape}"

    # Run loop
    input = torch.randn(2, 3, 10, device=device)
    h0 = torch.zeros(3, 20, device=device)
    with torch.no_grad():
        for i in range(input.size()[0]):
            h1 = gru_cell1(input[i], h0)
            h2 = gru_cell2(input[i], h0)

            # Make sure the final hidden states have the same shape
            assert h1.shape == h2.shape
            torch.testing.assert_close(h1, h2)
            h0 = h1
//  [python]
def test_multiagent_cnn(
        self,
        n_agents,
        centralised,
        share_params,
        batch,
        channels,
        x=15,
        y=15,
    ):
        torch.manual_seed(0)
        cnn = MultiAgentConvNet(
            n_agents=n_agents,
            centralised=centralised,
            share_params=share_params,
            in_features=channels,
            kernel_sizes=3,
        )
        if channels is None:
            channels = 3
        td = TensorDict(
            {
                "agents": TensorDict(
                    {"observation": torch.randn(*batch, n_agents, channels, x, y)},
                    [*batch, n_agents],
                )
            },
            batch_size=batch,
        )
        obs = td[("agents", "observation")]
        out = cnn(obs)
        assert out.shape[:-1] == (*batch, n_agents)
        if centralised and share_params:
            torch.testing.assert_allclose(out, out[..., :1, :].expand_as(out))
        else:
            for i in range(n_agents):
                for j in range(i + 1, n_agents):
                    assert not torch.allclose(out[..., i, :], out[..., j, :])
        obs[..., 0, 0, 0, 0] += 1
        out2 = cnn(obs)
        if centralised:
            # a modification to the input of agent 0 will impact all agents
            assert not torch.isclose(out, out2).all()
        elif n_agents > 1:
            assert not torch.isclose(out[..., 0, :], out2[..., 0, :]).all()
            torch.testing.assert_allclose(out[..., 1:, :], out2[..., 1:, :])

        obs = torch.randn(*batch, 1, channels, x, y).expand(
            *batch, n_agents, channels, x, y
        )
        out = cnn(obs)
        for i in range(n_agents):
            if share_params:
                # same input same output
                assert torch.allclose(out[..., i, :], out[..., 0, :])
            else:
                for j in range(i + 1, n_agents):
                    # same input different output
                    assert not torch.allclose(out[..., i, :], out[..., j, :])
//  [python]
def test_multiagent_cnn(
        self,
        n_agents,
        centralised,
        share_params,
        batch,
        channels,
        x=15,
        y=15,
    ):
        torch.manual_seed(0)
        cnn = MultiAgentConvNet(
            n_agents=n_agents,
            centralised=centralised,
            share_params=share_params,
            in_features=channels,
            kernel_sizes=3,
        )
        if channels is None:
            channels = 3
        td = TensorDict(
            {
                "agents": TensorDict(
                    {"observation": torch.randn(*batch, n_agents, channels, x, y)},
                    [*batch, n_agents],
                )
            },
            batch_size=batch,
        )
        obs = td[("agents", "observation")]
        out = cnn(obs)
        assert out.shape[:-1] == (*batch, n_agents)
        if centralised and share_params:
            torch.testing.assert_close(out, out[..., :1, :].expand_as(out))
        else:
            for i in range(n_agents):
                for j in range(i + 1, n_agents):
                    assert not torch.allclose(out[..., i, :], out[..., j, :])
        obs[..., 0, 0, 0, 0] += 1
        out2 = cnn(obs)
        if centralised:
            # a modification to the input of agent 0 will impact all agents
            assert not torch.isclose(out, out2).all()
        elif n_agents > 1:
            assert not torch.isclose(out[..., 0, :], out2[..., 0, :]).all()
            torch.testing.assert_close(out[..., 1:, :], out2[..., 1:, :])

        obs = torch.randn(*batch, 1, channels, x, y).expand(
            *batch, n_agents, channels, x, y
        )
        out = cnn(obs)
        for i in range(n_agents):
            if share_params:
                # same input same output
                assert torch.allclose(out[..., i, :], out[..., 0, :])
            else:
                for j in range(i + 1, n_agents):
                    # same input different output
                    assert not torch.allclose(out[..., i, :], out[..., j, :])
//  [python]
def test_python_lstm_cell(device, bias):
    lstm_cell1 = LSTMCell(10, 20, device=device, bias=bias)
    lstm_cell2 = nn.LSTMCell(10, 20, device=device, bias=bias)

    lstm_cell1.load_state_dict(lstm_cell2.state_dict())

    # Make sure parameters match
    for (k1, v1), (k2, v2) in zip(
        lstm_cell1.named_parameters(), lstm_cell2.named_parameters()
    ):
        assert k1 == k2, f"Parameter names do not match: {k1} != {k2}"
        assert (
            v1.shape == v2.shape
        ), f"Parameter shapes do not match: {k1} shape {v1.shape} != {k2} shape {v2.shape}"

    # Run loop
    input = torch.randn(2, 3, 10, device=device)
    h0 = torch.randn(3, 20, device=device)
    c0 = torch.randn(3, 20, device=device)
    with torch.no_grad():
        for i in range(input.size()[0]):
            h1, c1 = lstm_cell1(input[i], (h0, c0))
            h2, c2 = lstm_cell2(input[i], (h0, c0))

            # Make sure the final hidden states have the same shape
            assert h1.shape == h2.shape
            assert c1.shape == c2.shape
            torch.testing.assert_allclose(h1, h2)
            torch.testing.assert_allclose(c1, c2)
            h0 = h1
            c0 = c1
//  [python]
def test_python_lstm_cell(device, bias):
    lstm_cell1 = LSTMCell(10, 20, device=device, bias=bias)
    lstm_cell2 = nn.LSTMCell(10, 20, device=device, bias=bias)

    lstm_cell1.load_state_dict(lstm_cell2.state_dict())

    # Make sure parameters match
    for (k1, v1), (k2, v2) in zip(
        lstm_cell1.named_parameters(), lstm_cell2.named_parameters()
    ):
        assert k1 == k2, f"Parameter names do not match: {k1} != {k2}"
        assert (
            v1.shape == v2.shape
        ), f"Parameter shapes do not match: {k1} shape {v1.shape} != {k2} shape {v2.shape}"

    # Run loop
    input = torch.randn(2, 3, 10, device=device)
    h0 = torch.randn(3, 20, device=device)
    c0 = torch.randn(3, 20, device=device)
    with torch.no_grad():
        for i in range(input.size()[0]):
            h1, c1 = lstm_cell1(input[i], (h0, c0))
            h2, c2 = lstm_cell2(input[i], (h0, c0))

            # Make sure the final hidden states have the same shape
            assert h1.shape == h2.shape
            assert c1.shape == c2.shape
            torch.testing.assert_close(h1, h2)
            torch.testing.assert_close(c1, c2)
            h0 = h1
            c0 = c1
//  [python]
def test_non_tensor_input(self):
        def foo(x):
            a = x["a"]
            b = x["b"]
            return a + b

        x = {"a": torch.randn(2, 2), "b": torch.randn(2, 2)}
        ref_output = foo(x)
        func = make_fx(foo)
        gm = func(x)
        output = gm(x)
        torch.testing.assert_allclose(output, ref_output)
//  [python]
def test_non_tensor_input(self):
        def foo(x):
            a = x["a"]
            b = x["b"]
            return a + b

        x = {"a": torch.randn(2, 2), "b": torch.randn(2, 2)}
        ref_output = foo(x)
        func = make_fx(foo)
        gm = func(x)
        output = gm(x)
        torch.testing.assert_close(output, ref_output)
//  [python]
def test_python_gru(device, bias, dropout, batch_first, num_layers):
    B = 5
    T = 3
    gru1 = GRU(
        input_size=10,
        hidden_size=20,
        num_layers=num_layers,
        device=device,
        bias=bias,
        batch_first=batch_first,
    )
    gru2 = nn.GRU(
        input_size=10,
        hidden_size=20,
        num_layers=num_layers,
        device=device,
        bias=bias,
        batch_first=batch_first,
    )
    gru2.load_state_dict(gru1.state_dict())

    # Make sure parameters match
    for (k1, v1), (k2, v2) in zip(gru1.named_parameters(), gru2.named_parameters()):
        assert k1 == k2, f"Parameter names do not match: {k1} != {k2}"
        torch.testing.assert_allclose(v1, v2)
        assert (
            v1.shape == v2.shape
        ), f"Parameter shapes do not match: {k1} shape {v1.shape} != {k2} shape {v2.shape}"

    if batch_first:
        input = torch.randn(B, T, 10, device=device)
    else:
        input = torch.randn(T, B, 10, device=device)

    h0 = torch.randn(num_layers, 5, 20, device=device)

    # Test without hidden states
    with torch.no_grad():
        output1, h1 = gru1(input)
        output2, h2 = gru2(input)

    assert h1.shape == h2.shape
    assert output1.shape == output2.shape
    if dropout == 0.0:
        torch.testing.assert_allclose(output1, output2)
        torch.testing.assert_allclose(h1, h2)

    # Test with hidden states
    with torch.no_grad():
        output1, h1 = gru1(input, h0)
        output2, h2 = gru2(input, h0)

    assert h1.shape == h2.shape
    assert output1.shape == output2.shape
    if dropout == 0.0:
        torch.testing.assert_allclose(output1, output2)
        torch.testing.assert_allclose(h1, h2)
//  [python]
def test_python_gru(device, bias, dropout, batch_first, num_layers):
    B = 5
    T = 3
    gru1 = GRU(
        input_size=10,
        hidden_size=20,
        num_layers=num_layers,
        device=device,
        bias=bias,
        batch_first=batch_first,
    )
    gru2 = nn.GRU(
        input_size=10,
        hidden_size=20,
        num_layers=num_layers,
        device=device,
        bias=bias,
        batch_first=batch_first,
    )
    gru2.load_state_dict(gru1.state_dict())

    # Make sure parameters match
    for (k1, v1), (k2, v2) in zip(gru1.named_parameters(), gru2.named_parameters()):
        assert k1 == k2, f"Parameter names do not match: {k1} != {k2}"
        torch.testing.assert_close(v1, v2)
        assert (
            v1.shape == v2.shape
        ), f"Parameter shapes do not match: {k1} shape {v1.shape} != {k2} shape {v2.shape}"

    if batch_first:
        input = torch.randn(B, T, 10, device=device)
    else:
        input = torch.randn(T, B, 10, device=device)

    h0 = torch.randn(num_layers, 5, 20, device=device)

    # Test without hidden states
    with torch.no_grad():
        output1, h1 = gru1(input)
        output2, h2 = gru2(input)

    assert h1.shape == h2.shape
    assert output1.shape == output2.shape
    if dropout == 0.0:
        torch.testing.assert_close(output1, output2)
        torch.testing.assert_close(h1, h2)

    # Test with hidden states
    with torch.no_grad():
        output1, h1 = gru1(input, h0)
        output2, h2 = gru2(input, h0)

    assert h1.shape == h2.shape
    assert output1.shape == output2.shape
    if dropout == 0.0:
        torch.testing.assert_close(output1, output2)
        torch.testing.assert_close(h1, h2)