chore: update tests and dependencies

pyproject.toml

@@ -40,10 +40,12 @@ classifiers = [
 ]
 
 dependencies = [
-    "torch>=2.2",
-    "torch-geometric>=2.3.1,<2.5",
-    "einops>=0.6.1",
-    "anemoi-utils@git+https://github.com/ecmwf/anemoi-utils.git@feature/add-utils-for-training-and-models",
+    "torch==2.3",
+    "torch-geometric==2.4",
+    "einops==0.6.1",
+    "hydra-core==1.3",
+    "anemoi-datasets==0.2.1",
+    "anemoi-utils==0.1.9",
 ]
 
 [project.optional-dependencies]
@@ -53,8 +55,9 @@ docs = [
     # For building the documentation
 ]
 
-all = [
-]
+all = []
+
+tests = ["pytest", "hypothesis", "omegaconf"]
 
 dev = [
     "sphinx",

src/anemoi/models/data/data_indices/collection.py → src/anemoi/models/data_indices/collection.py

@@ -12,12 +12,12 @@
 import yaml
 from omegaconf import OmegaConf
 
-from anemoi.models.data.data_indices.index import BaseIndex
-from anemoi.models.data.data_indices.index import DataIndex
-from anemoi.models.data.data_indices.index import ModelIndex
-from anemoi.models.data.data_indices.tensor import BaseTensorIndex
-from anemoi.models.data.data_indices.tensor import InputTensorIndex
-from anemoi.models.data.data_indices.tensor import OutputTensorIndex
+from anemoi.models.data_indices.index import BaseIndex
+from anemoi.models.data_indices.index import DataIndex
+from anemoi.models.data_indices.index import ModelIndex
+from anemoi.models.data_indices.tensor import BaseTensorIndex
+from anemoi.models.data_indices.tensor import InputTensorIndex
+from anemoi.models.data_indices.tensor import OutputTensorIndex
 
 
 class IndexCollection:

src/anemoi/models/data/data_indices/index.py → src/anemoi/models/data_indices/index.py

@@ -7,8 +7,8 @@
 # nor does it submit to any jurisdiction.
 #
 
-from anemoi.models.data.data_indices.tensor import InputTensorIndex
-from anemoi.models.data.data_indices.tensor import OutputTensorIndex
+from anemoi.models.data_indices.tensor import InputTensorIndex
+from anemoi.models.data_indices.tensor import OutputTensorIndex
 
 
 class BaseIndex:

src/anemoi/models/layers/attention.py

@@ -24,8 +24,8 @@
 else:
     _FLASH_ATTENTION_AVAILABLE = True
 
-from src.anemoi.models.distributed.graph import shard_heads
-from src.anemoi.models.distributed.graph import shard_sequence
+from src.anemoi.models.distributed.transformer import shard_heads
+from src.anemoi.models.distributed.transformer import shard_sequence
 
 LOGGER = logging.getLogger(__name__)
 

src/anemoi/models/layers/block.py

@@ -25,10 +25,10 @@
 from anemoi.models.layers.conv import GraphConv
 from anemoi.models.layers.conv import GraphTransformerConv
 from anemoi.models.layers.mlp import MLP
-from src.anemoi.models.distributed.graph import shard_heads
-from src.anemoi.models.distributed.graph import shard_sequence
 from src.anemoi.models.distributed.graph import shard_tensor
 from src.anemoi.models.distributed.graph import sync_tensor
+from src.anemoi.models.distributed.transformer import shard_heads
+from src.anemoi.models.distributed.transformer import shard_sequence
 
 LOGGER = logging.getLogger(__name__)
 

src/anemoi/models/preprocessing/__init__.py

@@ -160,6 +160,8 @@ def forward(self, x, in_place: bool = True) -> Tensor:
 
     def _run_checks(self, x):
         """Run checks on the processed tensor."""
-        assert not torch.isnan(
-            x
-        ).any(), f"NaNs ({torch.isnan(x).sum()}) found in processed tensor after {self.__class__.__name__}."
+        if not self.inverse:
+            # Forward transformation checks:
+            assert not torch.isnan(
+                x
+            ).any(), f"NaNs ({torch.isnan(x).sum()}) found in processed tensor after {self.__class__.__name__}."

tests/data/data_indices/test_collection.py → tests/data_indices/test_collection.py

@@ -9,7 +9,7 @@
 import torch
 from omegaconf import DictConfig
 
-from anemoi.models.data.data_indices.collection import IndexCollection
+from anemoi.models.data_indices.collection import IndexCollection
 
 
 @pytest.fixture()

tests/data/data_indices/test_data_indices.py → tests/data_indices/test_data_indices.py

@@ -8,9 +8,9 @@
 import pytest
 import torch
 
-from anemoi.models.data.data_indices.index import DataIndex
-from anemoi.models.data.data_indices.tensor import InputTensorIndex
-from anemoi.models.data.data_indices.tensor import OutputTensorIndex
+from anemoi.models.data_indices.index import DataIndex
+from anemoi.models.data_indices.tensor import InputTensorIndex
+from anemoi.models.data_indices.tensor import OutputTensorIndex
 
 
 @pytest.fixture()

tests/layers/processor/test_graphconv_processor.py

@@ -88,7 +88,7 @@ def test_graphconv_processor_init(graphconv_processor, graphconv_init):
         _dst_grid_size,
         _trainable_size,
     ) = graphconv_init
-    assert graphconv_processor.num_layers == num_layers
+    assert graphconv_processor.num_chunks == num_chunks
     assert graphconv_processor.num_channels == num_channels
     assert graphconv_processor.chunk_size == num_layers // num_chunks
     assert isinstance(graphconv_processor.trainable, TrainableTensor)