Please address the comments from this code review:

## Individual Comments

### Comment 1
<location> `src/tdhook/weights/adapters.py:42-50` </location>
<code_context>
+                    adapter_input = cache_proxy.resolve()
+                else:
+                    adapter_input = kwargs.pop(DIRECTION_TO_RETURN[kwargs["direction"]])
+                return adapter(adapter_input, **kwargs)
+
+            return callback
</code_context>

<issue_to_address>
**suggestion (bug_risk):** Passing all kwargs to the adapter may cause unexpected keyword argument errors.

If the adapter's forward method doesn't accept arbitrary kwargs, this could cause runtime errors. Filter kwargs to only those accepted by the adapter, or clearly document the required adapter signature.

```suggestion
            import inspect

            def callback(**kwargs):
                nonlocal adapter, cache_proxy
                if cache_proxy is not None:
                    adapter_input = cache_proxy.resolve()
                else:
                    adapter_input = kwargs.pop(DIRECTION_TO_RETURN[kwargs["direction"]])
                # Filter kwargs to only those accepted by the adapter
                adapter_params = inspect.signature(adapter).parameters
                filtered_kwargs = {k: v for k, v in kwargs.items() if k in adapter_params}
                return adapter(adapter_input, **filtered_kwargs)

            return callback
```
</issue_to_address>

### Comment 2
<location> `tests/weights/test_pruning.py:79-88` </location>
<code_context>
+                torch.tensor(0.9),
+            )
+
+    def test_no_pruning(self):
+        model = nn.Linear(10, 10)
+        inp = torch.randn(10)
+        original_output = model(inp)
+
+        pruning = Pruning(importance_callback=_importance_cb, amount_to_prune=0)
+        ctx = pruning.prepare(model)
+        with ctx as hooked:
+            output = hooked(inp)
+            assert torch.allclose(output, original_output)
+
+    def test_skip_importance_cb(self):
</code_context>

<issue_to_address>
**suggestion (testing):** Add a test for negative pruning amounts.

Please add a test case for negative amount_to_prune to verify that the code correctly raises an error or handles the input as intended.

Suggested implementation:

```python
    def test_negative_pruning_amount(self):
        model = nn.Linear(10, 10)
        with pytest.raises(ValueError):
            Pruning(importance_callback=_importance_cb, amount_to_prune=-0.1)

    def test_skip_importance_cb(self):
        model = nn.Linear(10, 10)
        original_state = (model.weight.clone(), model.bias.clone())

```

Make sure that `pytest` is imported at the top of the file if it is not already:
```python
import pytest
```
Also, ensure that the `Pruning` class raises a `ValueError` for negative `amount_to_prune`. If it does not, you will need to add this validation in the `Pruning` class implementation.
</issue_to_address>

### Comment 3
<location> `tests/weights/test_pruning.py:90-98` </location>
<code_context>
+            output = hooked(inp)
+            assert torch.allclose(output, original_output)
+
+    def test_skip_importance_cb(self):
+        model = nn.Linear(10, 10)
+        original_state = (model.weight.clone(), model.bias.clone())
+
+        pruning = Pruning(importance_callback=_importance_cb_skip_weight, amount_to_prune=0.5)
+        ctx = pruning.prepare(model)
+        with ctx:
+            assert torch.allclose(model.weight, original_state[0])
+            assert not torch.allclose(model.bias, original_state[1])
</code_context>

<issue_to_address>
**suggestion (testing):** Consider testing skip_modules logic.

Please add a test for skip_modules by passing a custom function and confirming that the intended modules are excluded from pruning.
</issue_to_address>

### Comment 4
<location> `tests/weights/test_adapters.py:30-28` </location>
<code_context>
+        restored_out = default_test_model(data["input"])
+        assert torch.allclose(baseline_out, restored_out)
+
+    def test_adapter_crosslayer(self, default_test_model):
+        data = TensorDict({"input": torch.randn(4, 10)}, batch_size=4)
+        baseline_out = default_test_model(data["input"]).detach().clone()
+
+        adapters = {"linear2": (_DoubleAdapter(), "linear1", "linear2")}
+        ctx_factory = Adapters(adapters=adapters)
+
+        with ctx_factory.prepare(default_test_model) as hooked:
+            patched_data = hooked(data.clone())
+            patched_out = patched_data["output"]
+            assert not torch.allclose(baseline_out, patched_out)
+
+        restored_out = default_test_model(data["input"])
+        assert torch.allclose(baseline_out, restored_out)
</code_context>

<issue_to_address>
**suggestion (testing):** Consider adding tests for multiple adapters and directionality.

Expanding tests to cover multiple adapters, various directions, and edge cases like mismatched input/output shapes will improve test coverage and reliability.
</issue_to_address>

### Comment 5
<location> `tests/weights/test_pruning.py:15-17` </location>
<code_context>
def _importance_cb_skip_weight(parameter, parameter_name, **_):
    if parameter_name == "weight":
        return None
    return parameter

</code_context>

<issue_to_address>
**suggestion (code-quality):** We've found these issues:

- Lift code into else after jump in control flow ([`reintroduce-else`](https://docs.sourcery.ai/Reference/Default-Rules/refactorings/reintroduce-else/))
- Replace if statement with if expression ([`assign-if-exp`](https://docs.sourcery.ai/Reference/Default-Rules/refactorings/assign-if-exp/))

```suggestion
    return None if parameter_name == "weight" else parameter
```
</issue_to_address>

### Comment 6
<location> `tests/weights/test_pruning.py:21-23` </location>
<code_context>
def _importance_cb_skip_bias(parameter, parameter_name, **_):
    if parameter_name == "bias":
        return None
    return parameter

</code_context>

<issue_to_address>
**suggestion (code-quality):** We've found these issues:

- Lift code into else after jump in control flow ([`reintroduce-else`](https://docs.sourcery.ai/Reference/Default-Rules/refactorings/reintroduce-else/))
- Replace if statement with if expression ([`assign-if-exp`](https://docs.sourcery.ai/Reference/Default-Rules/refactorings/assign-if-exp/))

```suggestion
    return None if parameter_name == "bias" else parameter
```
</issue_to_address>