Support PyTorch 2.6

Author: chezou

poetry.lock

@@ -86,6 +86,7 @@ nmslib-metabrainz = {version = "^2.1.3", python = ">=3.11, <3.13", optional = tr
 torch = [
     {version = ">=1.6.0, <2.3.0", python = "<3.13", markers = "sys_platform == 'darwin' and platform_machine == 'x86_64'", optional = true},
     {version = ">=1.6.0, <3.0.0", python = "<3.13", optional = true},
+    {version = ">=2.6.0, <3.0.0", python = ">=3.13", optional = true},
 ]
 pytorch-lightning = {version = ">=1.6.0, <3.0.0", python = "<3.13", optional = true}
 

pyproject.toml

@@ -13,28 +13,15 @@
 #  limitations under the License.
 
 # pylint: disable=attribute-defined-outside-init,consider-using-enumerate
-import sys
-
 import numpy as np
 import pandas as pd
 import pytest
-
-try:
-    import torch
-except ImportError:
-    pass
-
+import torch
 from scipy import sparse
 
 from rectools.columns import Columns
 from rectools.dataset import Dataset
-
-try:
-    from rectools.dataset.torch_datasets import DSSMItemDataset, DSSMTrainDataset, DSSMUserDataset
-except ModuleNotFoundError:
-    pass
-
-pytestmark = pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+from rectools.dataset.torch_datasets import DSSMItemDataset, DSSMTrainDataset, DSSMUserDataset
 
 
 class WithFixtures:

tests/dataset/test_torch_dataset.py

@@ -50,7 +50,9 @@ def filter_warnings_decorator(func):  # type: ignore
 
 from ..data import INTERACTIONS
 
-pytestmark = pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+pytestmark = pytest.mark.skipif(
+    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
+)
 
 
 @filter_warnings_decorator

tests/models/nn/test_dssm.py

@@ -18,35 +18,29 @@
 import numpy as np
 import pandas as pd
 import pytest
+import torch
 
 try:
-    import torch
     from pytorch_lightning import seed_everything
 except ImportError:
     pass
 
 from rectools.columns import Columns
 from rectools.dataset import Dataset
 from rectools.dataset.dataset import DatasetSchema, EntitySchema
-
-try:
-    from rectools.models.nn.item_net import (
-        CatFeaturesItemNet,
-        IdEmbeddingsItemNet,
-        ItemNetBase,
-        ItemNetConstructorBase,
-        SumOfEmbeddingsConstructor,
-    )
-except ImportError:
-    CatFeaturesItemNet = object  # type: ignore
-    IdEmbeddingsItemNet = object  # type: ignore
-    ItemNetBase = object  # type: ignore
-    ItemNetConstructorBase = object  # type: ignore
-    SumOfEmbeddingsConstructor = object  # type: ignore
+from rectools.models.nn.item_net import (
+    CatFeaturesItemNet,
+    IdEmbeddingsItemNet,
+    ItemNetBase,
+    ItemNetConstructorBase,
+    SumOfEmbeddingsConstructor,
+)
 
 from ..data import DATASET, INTERACTIONS
 
-pytestmark = pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+pytestmark = pytest.mark.skipif(
+    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
+)
 
 
 class TestIdEmbeddingsItemNet:

tests/models/nn/test_item_net.py

@@ -19,10 +19,10 @@
 
 import pandas as pd
 import pytest
+import torch
 from pytest import FixtureRequest
 
 try:
-    import torch
     from pytorch_lightning import Trainer, seed_everything
     from pytorch_lightning.loggers import CSVLogger
 
@@ -46,7 +46,9 @@
 except NameError:
     pass
 
-pytestmark = pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+pytestmark = pytest.mark.skipif(
+    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
+)
 
 
 class TestTransformerModelBase:

tests/models/nn/transformers/test_base.py

@@ -13,28 +13,19 @@
 #  limitations under the License.
 
 import sys
-import types
 import typing as tp
 from functools import partial
 
 import numpy as np
 import pandas as pd
 import pytest
+import torch
 
 try:
-    import torch
     from pytorch_lightning import Trainer, seed_everything
 except ImportError:
-    torch = types.ModuleType("torch")
-    torch.Tensor = object  # type: ignore
-    torch.float = object  # type: ignore
     Trainer = object  # type: ignore
 
-    def tensor(*args: tp.Any, **kwargs: tp.Any) -> tp.Any:
-        return object()
-
-    torch.tensor = tensor
-
 from rectools import ExternalIds
 from rectools.columns import Columns
 from rectools.dataset import Dataset
@@ -64,7 +55,7 @@ def tensor(*args: tp.Any, **kwargs: tp.Any) -> tp.Any:
     pass
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13")
 class TestBERT4RecModel:
     def setup_method(self) -> None:
         self._seed_everything()
@@ -642,7 +633,7 @@ def _collate_fn_train(
         )
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13")
 class TestBERT4RecDataPreparator:
 
     def setup_method(self) -> None:
@@ -822,7 +813,7 @@ def test_get_dataloader_val(
             assert torch.equal(value, val_batch[key])
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13")
 class TestBERT4RecModelConfiguration:
     def setup_method(self) -> None:
         self._seed_everything()

tests/models/nn/transformers/test_bert4rec.py

@@ -12,7 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import sys
 import typing as tp
 
 import numpy as np
@@ -22,15 +21,9 @@
 from rectools.columns import Columns
 from rectools.dataset import Dataset, IdMap, Interactions
 from rectools.dataset.features import DenseFeatures
-
-try:
-    from rectools.models.nn.transformers.data_preparator import SequenceDataset, TransformerDataPreparatorBase
-except ImportError:
-    TransformerDataPreparatorBase = object  # type: ignore
+from rectools.models.nn.transformers.data_preparator import SequenceDataset, TransformerDataPreparatorBase
 from tests.testing_utils import assert_feature_set_equal, assert_id_map_equal, assert_interactions_set_equal
 
-pytestmark = pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
-
 
 class TestSequenceDataset:
 

tests/models/nn/transformers/test_data_preparator.py

@@ -15,21 +15,17 @@
 # pylint: disable=too-many-lines
 
 import sys
-import types
 import typing as tp
 from functools import partial
 
 import numpy as np
 import pandas as pd
 import pytest
+import torch
 
 try:
-    import torch
     from pytorch_lightning import Trainer, seed_everything
 except ImportError:
-    torch = types.ModuleType("torch")
-    torch.tensor = lambda x: None  # type: ignore
-    torch.Tensor = object  # type: ignore
     Trainer = object  # type: ignore
 
 from rectools import ExternalIds
@@ -61,7 +57,9 @@
 except NameError:
     pass
 
-pytestmark = pytest.mark.skipif(sys.version_info >= (3, 13), reason="`torch` is not compatible with Python >= 3.13")
+pytestmark = pytest.mark.skipif(
+    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
+)
 
 
 class TestSASRecModel: