Fix ruff errors

Author: Tim Joseph

src/tensorcontainer/distributions/sampling.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 import torch
 from torch.distributions import Distribution
 
@@ -47,7 +45,6 @@ def entropy(self):
         samples = self.base_dist.rsample((self.n,))
         logprob = self.base_dist.log_prob(samples)
         return -logprob.mean(0)
-        
+
     def log_prob(self, value):
         return self.base_dist.log_prob(value)
-

tests/tensor_dataclass/test_detach.py

@@ -34,8 +34,8 @@ def test_basic_detach(self, simple_tensor_data_instance, compile_mode):
         assert not detached_instance.b.requires_grad
 
         # Check that original tensors still have gradients (should be False for fixture)
-        assert test_instance.a.requires_grad == False
-        assert test_instance.b.requires_grad == False
+        assert not test_instance.a.requires_grad
+        assert not test_instance.b.requires_grad
 
         # Check that data is preserved
         assert torch.equal(detached_instance.a, test_instance.a)

tests/tensor_dict/test_recompilations.py

@@ -1,12 +1,11 @@
 import pytest
 import torch
-
-from tensorcontainer.tensor_dict import TensorDict
-from tests.conftest import skipif_no_compile
 import torch._dynamo
 import torch._dynamo.utils
 from torch._dynamo.testing import CompileCounter  # New import
 
+from tensorcontainer.tensor_dict import TensorDict
+from tests.conftest import skipif_no_compile
 
 keys = ["a", "b", "c", "d", "e", "f", "g"]
 
@@ -26,12 +25,10 @@ def test_getitem_recompilation(key):
     torch._dynamo.utils.counters.clear()
 
     # --- Sanity check with a simple lambda using CompileCounter ---
-    simple_lambda = lambda x: x + 1
-
     lambda_compile_counter = CompileCounter()
     # Using fullgraph=True as it was in the original attempts
     compiled_lambda = torch.compile(
-        simple_lambda, backend=lambda_compile_counter, fullgraph=True
+        lambda x: x + 1, backend=lambda_compile_counter, fullgraph=True
     )
 
     _ = compiled_lambda(torch.randn(1))

tests/tensor_dict/test_stack.py

@@ -3,10 +3,10 @@
 from torch._dynamo import exc as dynamo_exc
 
 from tensorcontainer.tensor_dict import TensorDict
+from tests.compile_utils import run_and_compare_compiled
 from tests.conftest import skipif_no_compile
 from tests.tensor_dict import common
 from tests.tensor_dict.common import compare_nested_dict, compute_stack_shape
-from tests.compile_utils import run_and_compare_compiled
 
 
 @pytest.fixture(autouse=True)
@@ -143,7 +143,7 @@ def stack_operation(tensor_dict_instance, stack_dimension):
         )
 
     compiled_stack_op = torch.compile(stack_operation, fullgraph=True)
-    with pytest.raises(dynamo_exc.Unsupported) as excinfo:
+    with pytest.raises(dynamo_exc.Unsupported):
         compiled_stack_op(td, dim)
 
 

tests/test_unsafe_construction.py

@@ -138,7 +138,7 @@ def test_unsafe_construction_nested_contexts(sample_tensors):
     """Test that nested unsafe construction contexts work correctly."""
     with TensorContainer.unsafe_construction():
         # Should work in outer context
-        container1 = SampleTensorDataClass(
+        SampleTensorDataClass(
             features=sample_tensors["incompatible_features"],
             labels=sample_tensors["labels"],
             shape=(4,),
@@ -210,7 +210,7 @@ def test_thread(thread_id, use_unsafe):
                 with TensorContainer.unsafe_construction():
                     # Add small delay to test concurrency
                     time.sleep(0.01)
-                    container = SampleTensorDataClass(
+                    SampleTensorDataClass(
                         features=sample_tensors["incompatible_features"],
                         labels=sample_tensors["incompatible_labels"],
                         shape=(4,),
@@ -219,7 +219,7 @@ def test_thread(thread_id, use_unsafe):
                     results[thread_id] = "success"
             else:
                 try:
-                    container = SampleTensorDataClass(
+                    SampleTensorDataClass(
                         features=sample_tensors["incompatible_features"],
                         labels=sample_tensors["incompatible_labels"],
                         shape=(4,),