lint

Author: vmoens

test/test_compile.py

@@ -14,6 +14,8 @@
 import pytest
 
 import torch
+
+from _utils_internal import is_npu_available
 from packaging import version
 
 from tensordict import (
@@ -39,8 +41,6 @@
 
 from tensordict.tensorclass import TensorClass
 
-from _utils_internal import is_npu_available
-
 from torch.utils._pytree import SUPPORTED_NODES, tree_map
 
 TORCH_VERSION = version.parse(version.parse(torch.__version__).base_version)
@@ -300,7 +300,8 @@ def test_to_device(td):
         assert td_device_c.device == torch.device(device)
 
     @pytest.mark.skipif(
-        is_npu_available(), reason="torch.device in torch.compile is not supported on NPU currently."
+        is_npu_available(),
+        reason="torch.device in torch.compile is not supported on NPU currently.",
     )
     def test_lock(self, mode):
         def locked_op(td):
@@ -593,7 +594,8 @@ def test_to_device(tc):
         assert tc_device_c.device == torch.device(device)
 
     @pytest.mark.skipif(
-        is_npu_available(), reason="torch.device in torch.compile is not supported on NPU currently."
+        is_npu_available(),
+        reason="torch.device in torch.compile is not supported on NPU currently.",
     )
     def test_tc_lock(self, mode):
         def locked_op(tc):

test/test_distributed.py

@@ -11,13 +11,13 @@
 import pytest
 import torch
 from _pytest.fixtures import fixture
+from _utils_internal import is_npu_available
 from packaging import version
 
 from packaging.version import parse
 
 from tensordict import LazyStackedTensorDict, MemoryMappedTensor, TensorDict
 from tensordict.utils import logger as tdlogger
-from _utils_internal import is_npu_available
 from torch import distributed as dist, multiprocessing as mp, nn
 from torch.distributed._tensor import (
     DeviceMesh,
@@ -109,7 +109,8 @@ def test_fsdp_module(self, tmpdir):
 
 
 @pytest.mark.skipif(
-    not is_npu_available() or not torch.npu.device_count() > 2, reason="not enough npu devices"
+    not is_npu_available() or not torch.npu.device_count() > 2,
+    reason="not enough npu devices",
 )
 class TestNPUFSDP:
     class MyDModule(nn.Module):
@@ -127,9 +128,7 @@ def forward(self, input):
     @classmethod
     def make_module(cls, device=None):
         with (
-            torch.device(f"npu:{device}")
-            if device is not None
-            else torch.device("npu")
+            torch.device(f"npu:{device}") if device is not None else torch.device("npu")
         ):
             my_module = cls.MyDModule()
             my_sharded_module = FSDP(my_module, device_id=device)
@@ -307,8 +306,8 @@ def server(queue):
                 },
                 [2],
             )
-                .expand(1, 2)
-                .contiguous()
+            .expand(1, 2)
+            .contiguous()
         )
         td.gather_and_stack(0)
         assert (td != 0).all()
@@ -380,8 +379,8 @@ def server(queue, op, async_op, return_premature):
                 },
                 [2],
             )
-                .expand(1, 2)
-                .contiguous()
+            .expand(1, 2)
+            .contiguous()
         )
         out = td.reduce(0, op=op, async_op=async_op, return_premature=return_premature)
         if not async_op:
@@ -864,8 +863,8 @@ def server(cls, queue):
                 },
                 [2],
             )
-                .expand(1, 2)
-                .contiguous()
+            .expand(1, 2)
+            .contiguous()
         )
         td.init_remote(dst=1)
 

test/test_nn.py

@@ -17,6 +17,7 @@
 
 import pytest
 import torch
+from _utils_internal import is_npu_available
 
 from tensordict import (
     is_tensor_collection,
@@ -57,7 +58,6 @@
     skip_existing,
 )
 from tensordict.tensorclass import TensorClass
-from _utils_internal import is_npu_available
 
 from torch import distributions, nn
 from torch.distributions import Categorical, Normal
@@ -2186,6 +2186,7 @@ def test_module_buffer():
 )
 def test_to_context(original_device, new_device, tc):
     if tc:
+
         class MyTC(TensorClass):
             x: torch.Tensor
             y: torch.Tensor | None = None

test/test_tensorclass.py

@@ -26,6 +26,7 @@
 import pytest
 import tensordict.utils
 import torch
+from _utils_internal import is_npu_available
 
 from tensordict import (
     assert_allclose_td,
@@ -45,7 +46,6 @@
 from tensordict._td import lazy_stack
 from tensordict.base import _GENERIC_NESTED_ERR
 from tensordict.tensorclass import from_dataclass
-from _utils_internal import is_npu_available
 
 from torch import Tensor
 
@@ -739,6 +739,7 @@ def test_disallowed_attributes(self):
             AttributeError,
             match="Attribute name reshape can't be used with @tensorclass",
         ):
+
             @tensorclass
             class MyInvalidClass:
                 x: torch.Tensor
@@ -1100,6 +1101,7 @@ class MyDataParent:
     @pytest.mark.parametrize("list_to_stack", [True, False])
     def test_indexing(self, list_to_stack):
         with set_list_to_stack(list_to_stack):
+
             @tensorclass
             class MyDataNested:
                 X: torch.Tensor
@@ -1436,8 +1438,8 @@ class MyDataNested(TensorClass):
         assert (
             repeated.X
             == X.repeat_interleave(
-            torch.tensor([2, 3, 4, 5], device=data.device), dim=1
-        )
+                torch.tensor([2, 3, 4, 5], device=data.device), dim=1
+            )
         ).all()
 
     def test_reshape(self):
@@ -2888,20 +2890,23 @@ class FuncAutoCast:
 class TestShadow:
     def test_no_shadow(self):
         with pytest.raises(AttributeError):
+
             @tensorclass
             class MyClass:
                 x: str
                 y: int
                 batch_size: Any
 
         with pytest.raises(AttributeError):
+
             @tensorclass
             class MyClass:  # noqa: F811
                 x: str
                 y: int
                 names: Any
 
         with pytest.raises(AttributeError):
+
             @tensorclass
             class MyClass:  # noqa: F811
                 x: str
@@ -3099,7 +3104,7 @@ class MyClass:
         _ = c / 1
         _ = 1 / c
 
-        _ = c ** 1
+        _ = c**1
         # not implemented
         # 1 ** c
 
@@ -3299,13 +3304,15 @@ class TensorOnly:
             c: torch.Tensor | None = None
 
         with pytest.raises(TypeError, match="tensor_only"):
+
             @tensorclass(tensor_only=True, nocast=True)
             class TensorOnlyNocast:
                 a: torch.Tensor
                 b: torch.Tensor
                 c: torch.Tensor | None = None
 
         with pytest.raises(TypeError, match="tensor_only"):
+
             @tensorclass(tensor_only=True, autocast=True)
             class TensorOnlyAutocast:
                 a: torch.Tensor
@@ -3330,6 +3337,7 @@ class TensorOnly(TensorClass["tensor_only"]):
             TypeError,
             match="tensor_only requires types to be Tensor, Tensor-subtrypes or None",
         ):
+
             class TensorOnlyAny(TensorClass["tensor_only"]):
                 a: torch.Tensor
                 b: Any
@@ -3339,6 +3347,7 @@ class TensorOnlyAny(TensorClass["tensor_only"]):
             TypeError,
             match="tensor_only requires types to be Tensor, Tensor-subtrypes or None",
         ):
+
             class TensorOnlyStr(TensorClass["tensor_only"]):
                 a: torch.Tensor
                 b: torch.Tensor | str
@@ -3348,6 +3357,7 @@ class TensorOnlyStr(TensorClass["tensor_only"]):
             TypeError,
             match="tensor_only requires types to be Tensor, Tensor-subtrypes or None",
         ):
+
             class TensorOnlyStrUnion(TensorClass["tensor_only"]):
                 a: torch.Tensor
                 b: torch.Tensor