From f96bcf9ea4f76718ed7952c8d269f93cb759b099 Mon Sep 17 00:00:00 2001
From: Jan Polster <jan.polster@ecmwf.int>
Date: Wed, 2 Oct 2024 15:00:05 +0200
Subject: [PATCH 01/18] Feature/mapper chunking (#46)

* feat: add chunking strategy for GraphTransformerMapperBlock to reduce peak memory usage during inference

* test: add tests for mapper chunking

* docs: update changelog

* docs: document new ANEMOI_INFERENCE_NUM_CHUNKS environment variable
---
 CHANGELOG.md                                  |  2 +
 docs/modules/layers.rst                       | 14 ++++
 src/anemoi/models/distributed/khop_edges.py   | 69 +++++++++++++------
 src/anemoi/models/layers/block.py             | 47 +++++++++++--
 src/anemoi/models/layers/mapper.py            |  4 +-
 src/anemoi/models/layers/processor.py         |  4 +-
 .../block/test_block_graphtransformer.py      | 47 ++++++++++++-
 .../test_preprocessor_normalizer.py           |  2 +-
 8 files changed, 158 insertions(+), 31 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 68106af0..4798ac9a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,12 +11,14 @@ Keep it human-readable, your future self will thank you!
 ## [Unreleased](https://github.com/ecmwf/anemoi-models/compare/0.3.0...HEAD)
 
 ### Added
+
 - Codeowners file
 - Pygrep precommit hooks
 - Docsig precommit hooks
 - Changelog merge strategy
 - configurabilty of the dropout probability in the the MultiHeadSelfAttention module
 - Variable Bounding as configurable model layers [#13](https://github.com/ecmwf/anemoi-models/issues/13)
+- GraphTransformerMapperBlock chunking to reduce memory usage during inference [#46](https://github.com/ecmwf/anemoi-models/pull/46)
 
 ### Changed
 - Bugfixes for CI
diff --git a/docs/modules/layers.rst b/docs/modules/layers.rst
index ff21b5ab..8082e47c 100644
--- a/docs/modules/layers.rst
+++ b/docs/modules/layers.rst
@@ -2,6 +2,20 @@
  Layers
 ########
 
+***********************
+ Environment Variables
+***********************
+
+``ANEMOI_INFERENCE_NUM_CHUNKS``
+===============================
+
+This environment variable controls the number of chunks used in the
+`Mapper` during inference. Setting this variable allows the model to
+split large computations into a specified number of smaller chunks,
+reducing memory overhead. If not set, it falls back to the default value
+of, 1 i.e. no chunking. See pull request `#46
+<https://github.com/ecmwf/anemoi-models/pull/46>`_.
+
 *********
  Mappers
 *********
diff --git a/src/anemoi/models/distributed/khop_edges.py b/src/anemoi/models/distributed/khop_edges.py
index 5b4bd815..73397d9a 100644
--- a/src/anemoi/models/distributed/khop_edges.py
+++ b/src/anemoi/models/distributed/khop_edges.py
@@ -46,7 +46,7 @@ def get_k_hop_edges(nodes: Tensor, edge_attr: Tensor, edge_index: Adj, num_hops:
     return edge_attr[mask_to_index(edge_mask_k)], edge_index_k
 
 
-def sort_edges_1hop(
+def sort_edges_1hop_sharding(
     num_nodes: Union[int, tuple[int, int]],
     edge_attr: Tensor,
     edge_index: Adj,
@@ -74,29 +74,56 @@ def sort_edges_1hop(
     if mgroup:
         num_chunks = dist.get_world_size(group=mgroup)
 
-        if isinstance(num_nodes, int):
-            node_chunks = torch.arange(num_nodes, device=edge_index.device).tensor_split(num_chunks)
-        else:
-            nodes_src = torch.arange(num_nodes[0], device=edge_index.device)
-            node_chunks = torch.arange(num_nodes[1], device=edge_index.device).tensor_split(num_chunks)
-
-        edge_index_list = []
-        edge_attr_list = []
-        for node_chunk in node_chunks:
-            if isinstance(num_nodes, int):
-                edge_attr_chunk, edge_index_chunk = get_k_hop_edges(node_chunk, edge_attr, edge_index)
-            else:
-                edge_index_chunk, edge_attr_chunk = bipartite_subgraph(
-                    (nodes_src, node_chunk),
-                    edge_index,
-                    edge_attr,
-                    size=(num_nodes[0], num_nodes[1]),
-                )
-            edge_index_list.append(edge_index_chunk)
-            edge_attr_list.append(edge_attr_chunk)
+        edge_attr_list, edge_index_list = sort_edges_1hop_chunks(num_nodes, edge_attr, edge_index, num_chunks)
+
         edge_index_shapes = [x.shape for x in edge_index_list]
         edge_attr_shapes = [x.shape for x in edge_attr_list]
 
         return torch.cat(edge_attr_list, dim=0), torch.cat(edge_index_list, dim=1), edge_attr_shapes, edge_index_shapes
 
     return edge_attr, edge_index, [], []
+
+
+def sort_edges_1hop_chunks(
+    num_nodes: Union[int, tuple[int, int]], edge_attr: Tensor, edge_index: Adj, num_chunks: int
+) -> tuple[list[Tensor], list[Adj]]:
+    """Rearanges edges into 1 hop neighbourhood chunks.
+
+    Parameters
+    ----------
+    num_nodes : Union[int, tuple[int, int]]
+        Number of (target) nodes in Graph, tuple for bipartite graph
+    edge_attr : Tensor
+        edge attributes
+    edge_index : Adj
+        edge index
+    num_chunks : int
+        number of chunks used if mgroup is None
+
+    Returns
+    -------
+    tuple[list[Tensor], list[Adj]]
+        list of sorted edge attribute chunks, list of sorted edge_index chunks
+    """
+    if isinstance(num_nodes, int):
+        node_chunks = torch.arange(num_nodes, device=edge_index.device).tensor_split(num_chunks)
+    else:
+        nodes_src = torch.arange(num_nodes[0], device=edge_index.device)
+        node_chunks = torch.arange(num_nodes[1], device=edge_index.device).tensor_split(num_chunks)
+
+    edge_index_list = []
+    edge_attr_list = []
+    for node_chunk in node_chunks:
+        if isinstance(num_nodes, int):
+            edge_attr_chunk, edge_index_chunk = get_k_hop_edges(node_chunk, edge_attr, edge_index)
+        else:
+            edge_index_chunk, edge_attr_chunk = bipartite_subgraph(
+                (nodes_src, node_chunk),
+                edge_index,
+                edge_attr,
+                size=(num_nodes[0], num_nodes[1]),
+            )
+        edge_index_list.append(edge_index_chunk)
+        edge_attr_list.append(edge_attr_chunk)
+
+    return edge_attr_list, edge_index_list
diff --git a/src/anemoi/models/layers/block.py b/src/anemoi/models/layers/block.py
index 7fd3627c..129e22ed 100644
--- a/src/anemoi/models/layers/block.py
+++ b/src/anemoi/models/layers/block.py
@@ -8,6 +8,7 @@
 #
 
 import logging
+import os
 from abc import ABC
 from abc import abstractmethod
 from typing import Optional
@@ -23,6 +24,7 @@
 
 from anemoi.models.distributed.graph import shard_tensor
 from anemoi.models.distributed.graph import sync_tensor
+from anemoi.models.distributed.khop_edges import sort_edges_1hop_chunks
 from anemoi.models.distributed.transformer import shard_heads
 from anemoi.models.distributed.transformer import shard_sequence
 from anemoi.models.layers.attention import MultiHeadSelfAttention
@@ -32,6 +34,9 @@
 
 LOGGER = logging.getLogger(__name__)
 
+# Number of Mapper chunks used during inference (https://github.com/ecmwf/anemoi-models/pull/46)
+NUM_CHUNKS_INFERENCE = int(os.environ.get("ANEMOI_INFERENCE_NUM_CHUNKS", "1"))
+
 
 class BaseBlock(nn.Module, ABC):
     """Base class for network blocks."""
@@ -498,14 +503,48 @@ def forward(
             ), "Only batch size of 1 is supported when model is sharded across GPUs"
 
         query, key, value, edges = self.shard_qkve_heads(query, key, value, edges, shapes, batch_size, model_comm_group)
-        out = self.conv(query=query, key=key, value=value, edge_attr=edges, edge_index=edge_index, size=size)
+
+        num_chunks = self.num_chunks if self.training else NUM_CHUNKS_INFERENCE
+
+        if num_chunks > 1:
+            # split 1-hop edges into chunks, compute self.conv chunk-wise and aggregate
+            edge_attr_list, edge_index_list = sort_edges_1hop_chunks(
+                num_nodes=size, edge_attr=edges, edge_index=edge_index, num_chunks=num_chunks
+            )
+            for i in range(num_chunks):
+                out1 = self.conv(
+                    query=query,
+                    key=key,
+                    value=value,
+                    edge_attr=edge_attr_list[i],
+                    edge_index=edge_index_list[i],
+                    size=size,
+                )
+                if i == 0:
+                    out = torch.zeros_like(out1, device=out1.device)
+                out = out + out1
+        else:
+            out = self.conv(query=query, key=key, value=value, edge_attr=edges, edge_index=edge_index, size=size)
+
         out = self.shard_output_seq(out, shapes, batch_size, model_comm_group)
-        out = self.projection(out + x_r)
+
+        # compute out = self.projection(out + x_r) in chunks:
+        out = torch.cat([self.projection(chunk) for chunk in torch.tensor_split(out + x_r, num_chunks, dim=0)], dim=0)
 
         out = out + x_skip[1]
-        nodes_new_dst = self.node_dst_mlp(out) + out
 
-        nodes_new_src = self.node_src_mlp(x_skip[0]) + x_skip[0] if self.update_src_nodes else x_skip[0]
+        # compute nodes_new_dst = self.node_dst_mlp(out) + out in chunks:
+        nodes_new_dst = torch.cat(
+            [self.node_dst_mlp(chunk) + chunk for chunk in out.tensor_split(num_chunks, dim=0)], dim=0
+        )
+
+        if self.update_src_nodes:
+            # compute nodes_new_src = self.node_src_mlp(out) + out in chunks:
+            nodes_new_src = torch.cat(
+                [self.node_src_mlp(chunk) + chunk for chunk in x_skip[0].tensor_split(num_chunks, dim=0)], dim=0
+            )
+        else:
+            nodes_new_src = x_skip[0]
 
         nodes_new = (nodes_new_src, nodes_new_dst)
 
diff --git a/src/anemoi/models/layers/mapper.py b/src/anemoi/models/layers/mapper.py
index 09670410..dc5c1f08 100644
--- a/src/anemoi/models/layers/mapper.py
+++ b/src/anemoi/models/layers/mapper.py
@@ -23,7 +23,7 @@
 
 from anemoi.models.distributed.graph import gather_tensor
 from anemoi.models.distributed.graph import shard_tensor
-from anemoi.models.distributed.khop_edges import sort_edges_1hop
+from anemoi.models.distributed.khop_edges import sort_edges_1hop_sharding
 from anemoi.models.distributed.shapes import change_channels_in_shape
 from anemoi.models.distributed.shapes import get_shape_shards
 from anemoi.models.layers.block import GraphConvMapperBlock
@@ -484,7 +484,7 @@ def __init__(
     def prepare_edges(self, size, batch_size, model_comm_group=None):
         edge_attr = self.trainable(self.edge_attr, batch_size)
         edge_index = self._expand_edges(self.edge_index_base, self.edge_inc, batch_size)
-        edge_attr, edge_index, shapes_edge_attr, shapes_edge_idx = sort_edges_1hop(
+        edge_attr, edge_index, shapes_edge_attr, shapes_edge_idx = sort_edges_1hop_sharding(
             size, edge_attr, edge_index, model_comm_group
         )
 
diff --git a/src/anemoi/models/layers/processor.py b/src/anemoi/models/layers/processor.py
index 6ba8eb1f..f35845ae 100644
--- a/src/anemoi/models/layers/processor.py
+++ b/src/anemoi/models/layers/processor.py
@@ -18,7 +18,7 @@
 from torch_geometric.data import HeteroData
 
 from anemoi.models.distributed.graph import shard_tensor
-from anemoi.models.distributed.khop_edges import sort_edges_1hop
+from anemoi.models.distributed.khop_edges import sort_edges_1hop_sharding
 from anemoi.models.distributed.shapes import change_channels_in_shape
 from anemoi.models.distributed.shapes import get_shape_shards
 from anemoi.models.layers.chunk import GNNProcessorChunk
@@ -235,7 +235,7 @@ def forward(
         edge_attr = self.trainable(self.edge_attr, batch_size)
         edge_index = self._expand_edges(self.edge_index_base, self.edge_inc, batch_size)
         target_nodes = sum(x[0] for x in shape_nodes)
-        edge_attr, edge_index, shapes_edge_attr, shapes_edge_idx = sort_edges_1hop(
+        edge_attr, edge_index, shapes_edge_attr, shapes_edge_idx = sort_edges_1hop_sharding(
             target_nodes,
             edge_attr,
             edge_index,
diff --git a/tests/layers/block/test_block_graphtransformer.py b/tests/layers/block/test_block_graphtransformer.py
index 00af258f..3d1c1da4 100644
--- a/tests/layers/block/test_block_graphtransformer.py
+++ b/tests/layers/block/test_block_graphtransformer.py
@@ -5,10 +5,13 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import importlib
+
 import pytest
 import torch
 import torch.nn as nn
 
+import anemoi.models.layers.block
 from anemoi.models.layers.block import GraphTransformerMapperBlock
 from anemoi.models.layers.block import GraphTransformerProcessorBlock
 from anemoi.models.layers.conv import GraphTransformerConv
@@ -303,9 +306,10 @@ def test_GraphTransformerMapperBlock_forward_backward(init, mapper_block):
     edge_index = torch.randint(1, 10, (2, 10))
     shapes = (10, 10, 10)
     batch_size = 1
+    size = (10, 10)
 
     # Forward pass
-    output, _ = block(x, edge_attr, edge_index, shapes, batch_size)
+    output, _ = block(x, edge_attr, edge_index, shapes, batch_size, size=size)
 
     # Check output shape
     assert output[0].shape == (10, out_channels)
@@ -327,3 +331,44 @@ def test_GraphTransformerMapperBlock_forward_backward(init, mapper_block):
         assert (
             param.grad.shape == param.shape
         ), f"param.grad.shape ({param.grad.shape}) != param.shape ({param.shape}) for {param}"
+
+
+def test_GraphTransformerMapperBlock_chunking(init, mapper_block, monkeypatch):
+    (
+        in_channels,
+        _hidden_dim,
+        _out_channels,
+        edge_dim,
+        _bias,
+        _activation,
+        _num_heads,
+        _num_chunks,
+    ) = init
+    # Initialize GraphTransformerMapperBlock
+    block = mapper_block
+
+    # Generate random input tensor
+    x = (torch.randn((10, in_channels)), torch.randn((10, in_channels)))
+    edge_attr = torch.randn((10, edge_dim))
+    edge_index = torch.randint(1, 10, (2, 10))
+    shapes = (10, 10, 10)
+    batch_size = 1
+    size = (10, 10)
+    num_chunks = torch.randint(2, 10, (1,)).item()
+
+    # manually set to non-training mode
+    block.eval()
+
+    # result with chunks
+    monkeypatch.setenv("ANEMOI_INFERENCE_NUM_CHUNKS", str(num_chunks))
+    importlib.reload(anemoi.models.layers.block)
+    out_chunked, _ = block(x, edge_attr, edge_index, shapes, batch_size, size=size)
+    # result without chunks, reload block for new env variable
+    monkeypatch.setenv("ANEMOI_INFERENCE_NUM_CHUNKS", "1")
+    importlib.reload(anemoi.models.layers.block)
+    out, _ = block(x, edge_attr, edge_index, shapes, batch_size, size=size)
+
+    assert out[0].shape == out_chunked[0].shape, f"out.shape ({out.shape}) != out_chunked.shape ({out_chunked.shape})"
+    assert out[1].shape == out_chunked[1].shape, f"out.shape ({out.shape}) != out_chunked.shape ({out_chunked.shape})"
+    assert torch.allclose(out[0], out_chunked[0], atol=1e-4), "out != out_chunked"
+    assert torch.allclose(out[1], out_chunked[1], atol=1e-4), "out != out_chunked"
diff --git a/tests/preprocessing/test_preprocessor_normalizer.py b/tests/preprocessing/test_preprocessor_normalizer.py
index 8056865d..3a2327ea 100644
--- a/tests/preprocessing/test_preprocessor_normalizer.py
+++ b/tests/preprocessing/test_preprocessor_normalizer.py
@@ -68,7 +68,7 @@ def remap_normalizer():
     }
     name_to_index = {"x": 0, "y": 1, "z": 2, "q": 3, "other": 4}
     data_indices = IndexCollection(config=config, name_to_index=name_to_index)
-    return InputNormalizer(config=config.data.normalizer, statistics=statistics, data_indices=data_indices)
+    return InputNormalizer(config=config.data.normalizer, data_indices=data_indices, statistics=statistics)
 
 
 def test_normalizer_not_inplace(input_normalizer) -> None:

From 116cc46e07f9fbdd0a24804203dd5cdc7bae7173 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 7 Oct 2024 18:41:04 +0000
Subject: [PATCH 02/18] [pre-commit.ci] pre-commit autoupdate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

updates:
- [github.com/pre-commit/pre-commit-hooks: v4.6.0 → v5.0.0](https://github.com/pre-commit/pre-commit-hooks/compare/v4.6.0...v5.0.0)
- [github.com/astral-sh/ruff-pre-commit: v0.6.4 → v0.6.9](https://github.com/astral-sh/ruff-pre-commit/compare/v0.6.4...v0.6.9)
- [github.com/tox-dev/pyproject-fmt: 2.2.3 → 2.2.4](https://github.com/tox-dev/pyproject-fmt/compare/2.2.3...2.2.4)
- [github.com/jshwi/docsig: v0.60.1 → v0.64.0](https://github.com/jshwi/docsig/compare/v0.60.1...v0.64.0)
---
 .pre-commit-config.yaml | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f3c39623..82ef0de1 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,12 +5,12 @@ repos:
   - id: clear-notebooks-output
     name: clear-notebooks-output
     files: tools/.*\.ipynb$
-    stages: [commit]
+    stages: [Nonepre-commitNone]
     language: python
     entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
     additional_dependencies: [jupyter]
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.6.0
+  rev: v5.0.0
   hooks:
   - id: check-yaml # Check YAML files for syntax errors only
     args: [--unsafe, --allow-multiple-documents]
@@ -40,7 +40,7 @@ repos:
     - --force-single-line-imports
     - --profile black
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.6.4
+  rev: v0.6.9
   hooks:
   - id: ruff
     # Next line if for documenation cod snippets
@@ -66,11 +66,11 @@ repos:
   - id: docconvert
     args: ["numpy"]
 - repo: https://github.com/tox-dev/pyproject-fmt
-  rev: "2.2.3"
+  rev: "2.2.4"
   hooks:
   - id: pyproject-fmt
 -   repo: https://github.com/jshwi/docsig # Check docstrings against function sig
-    rev: v0.60.1
+    rev: v0.64.0
     hooks:
     -   id: docsig
         args:

From 72d8f21aa4c1e4f55a42a116c6719b8dbdb4462f Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Thu, 10 Oct 2024 14:28:53 +0100
Subject: [PATCH 03/18] Add workflow to sync repos

---
 .github/workflows/ci.yml              |  4 ++--
 .github/workflows/push-to-private.yml | 33 +++++++++++++++++++++++++++
 .github/workflows/python-publish.yml  |  2 ++
 3 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/push-to-private.yml

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f089733b..25110667 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -35,7 +35,7 @@ jobs:
   # Run CI including downstream packages on self-hosted runners
   downstream-ci:
     name: downstream-ci
-    if: ${{ !github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci' }}
+    if: ${{ !contains(github.repository, 'private') && (!github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci') }}
     uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci.yml@main
     with:
       anemoi-models: ecmwf/anemoi-models@${{ github.event.pull_request.head.sha || github.sha }}
@@ -45,7 +45,7 @@ jobs:
    # Build downstream packages on HPC
   downstream-ci-hpc:
     name: downstream-ci-hpc
-    if: ${{ !github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci' }}
+    if: ${{ !contains(github.repository, 'private') && (!github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci') }}
     uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci-hpc.yml@main
     with:
       anemoi-models: ecmwf/anemoi-models@${{ github.event.pull_request.head.sha || github.sha }}
diff --git a/.github/workflows/push-to-private.yml b/.github/workflows/push-to-private.yml
new file mode 100644
index 00000000..4cc53efd
--- /dev/null
+++ b/.github/workflows/push-to-private.yml
@@ -0,0 +1,33 @@
+name: Push to private repository
+
+on:
+  push:
+    branches:
+      - develop
+
+jobs:
+  push_changes:
+    if: ${{ !contains(github.repository, 'private') }}
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout source repository
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+        fetch-tags: true
+
+    - name: Set up Git configuration
+      run: |
+        git config user.name "github-actions[bot]"
+        git config user.email "github-actions[bot]@users.noreply.github.com"
+
+    - name: Setup SSH key
+      uses: webfactory/ssh-agent@v0.5.0
+      with:
+        ssh-private-key: ${{ secrets.KEY_TO_PRIVATE }}
+
+    - name: Push changes to private repository
+      run: |
+        git remote add private git@github.com:${{ github.repository }}-private.git
+        git push --set-upstream private develop
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 2c28139b..37781907 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -9,11 +9,13 @@ on:
 
 jobs:
   quality:
+    if: ${{ !contains(github.repository, 'private') }}
     uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2
     with:
       skip-hooks: "no-commit-to-branch"
 
   checks:
+    if: ${{ !contains(github.repository, 'private') }}
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.x"]

From a8d19375f2ffc4062d05f16feeaf079ded093e13 Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Thu, 10 Oct 2024 14:30:36 +0100
Subject: [PATCH 04/18] Update changelog

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4798ac9a..2c99c52e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@ Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-models/compare/0.3.0...HEAD)
 
+- Add synchronisation workflow
+
 ### Added
 
 - Codeowners file

From 88d88dd7badac606a83065aa9cbf638c8a5184cd Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Thu, 10 Oct 2024 16:33:05 +0100
Subject: [PATCH 05/18] add link to transform

---
 CHANGELOG.md   | 1 +
 docs/conf.py   | 4 ++++
 docs/index.rst | 1 +
 3 files changed, 6 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4798ac9a..aec182d4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ Keep it human-readable, your future self will thank you!
 
 ### Added
 
+- Add anemoi-transform link to documentation
 - Codeowners file
 - Pygrep precommit hooks
 - Docsig precommit hooks
diff --git a/docs/conf.py b/docs/conf.py
index d7e0ff3f..8a5d5688 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -100,6 +100,10 @@
         "https://anemoi-registry.readthedocs.io/en/latest/",
         ("../../anemoi-registry/docs/_build/html/objects.inv", None),
     ),
+    "anemoi-transform": (
+        "https://anemoi-transform.readthedocs.io/en/latest/",
+        ("../../anemoi-transform/docs/_build/html/objects.inv", None),
+    ),
 }
 
 # -- Options for HTML output -------------------------------------------------
diff --git a/docs/index.rst b/docs/index.rst
index fa8f66b5..91da8bdb 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -48,6 +48,7 @@ weather forecasting.
 *****************
 
 -  :ref:`anemoi-utils <anemoi-utils:index-page>`
+-  :ref:`anemoi-transform <anemoi-transform:index-page>`
 -  :ref:`anemoi-datasets <anemoi-datasets:index-page>`
 -  :ref:`anemoi-models <anemoi-models:index-page>`
 -  :ref:`anemoi-graphs <anemoi-graphs:index-page>`

From 7b94e79aa2bdc8488bf2e5d61a171a639020d604 Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Tue, 15 Oct 2024 14:42:29 +0100
Subject: [PATCH 06/18] Update CODEOWNERS

---
 .github/CODEOWNERS | 6 +++---
 CHANGELOG.md       | 1 +
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index a2c619f5..74bdac0a 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1,6 +1,6 @@
 # CODEOWNERS file
 
 # Protect workflow files
-/.github/ @theissenhelen @jesperdramsch @gmertes
-/.pre-commit-config.yaml @theissenhelen @jesperdramsch @gmertes
-/pyproject.toml @theissenhelen @jesperdramsch @gmertes
+/.github/ @theissenhelen @jesperdramsch @gmertes @b8raoult @floriankrb @anaprietonem @HCookie @JPXKQX @mchantry
+/.pre-commit-config.yaml @theissenhelen @jesperdramsch @gmertes @b8raoult @floriankrb @anaprietonem @HCookie @JPXKQX @mchantry
+/pyproject.toml @theissenhelen @jesperdramsch @gmertes @b8raoult @floriankrb @anaprietonem @HCookie @JPXKQX @mchantry
diff --git a/CHANGELOG.md b/CHANGELOG.md
index aec182d4..20c26f08 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ Keep it human-readable, your future self will thank you!
 - Bugfixes for CI
 - Change Changelog CI to run after successful publish
 - pytest for downstream-ci-hpc
+- Update CODEOWNERS
 
 ### Removed
 

From b59d1ae0ef6b1e5efab9a8e987aeeead2023d428 Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Wed, 16 Oct 2024 04:14:38 +0000
Subject: [PATCH 07/18] Update push to pypi

---
 .github/workflows/python-publish.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 37781907..2f692eda 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -24,6 +24,7 @@ jobs:
       python-version: ${{ matrix.python-version }}
 
   deploy:
+    if: ${{ !contains(github.repository, 'private') }}
     needs: [checks, quality]
     uses: ecmwf-actions/reusable-workflows/.github/workflows/cd-pypi.yml@v2
     secrets: inherit

From 78428d0e224f4c17305ecd474882f4701a2c6536 Mon Sep 17 00:00:00 2001
From: b8raoult <53792887+b8raoult@users.noreply.github.com>
Date: Fri, 18 Oct 2024 07:57:40 +0100
Subject: [PATCH 08/18] Update .pre-commit-config.yaml (#59)

---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f3c39623..b929683e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -44,7 +44,7 @@ repos:
   hooks:
   - id: ruff
     # Next line if for documenation cod snippets
-    exclude: '^[^_].*_\.py$'
+    exclude: '.*/[^_].*_\.py$'
     args:
     - --line-length=120
     - --fix

From 6de3bdc531cb2677cca33737205ffd885cd15eab Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Fri, 18 Oct 2024 06:59:51 +0000
Subject: [PATCH 09/18] fix: typo

---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 82ef0de1..e01d6a37 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,7 +5,7 @@ repos:
   - id: clear-notebooks-output
     name: clear-notebooks-output
     files: tools/.*\.ipynb$
-    stages: [Nonepre-commitNone]
+    stages: [pre-commit]
     language: python
     entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
     additional_dependencies: [jupyter]

From 366fc4689e79f4790d4f04e18c6a0e5ba09cec4a Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Sun, 20 Oct 2024 19:33:05 +0000
Subject: [PATCH 10/18] Fix pre-commit regex

---
 .pre-commit-config.yaml | 3 +--
 CHANGELOG.md            | 1 +
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9cb0fcf3..4de59323 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -43,13 +43,12 @@ repos:
   rev: v0.6.9
   hooks:
   - id: ruff
-    # Next line if for documenation cod snippets
-    exclude: '.*/[^_].*_\.py$'
     args:
     - --line-length=120
     - --fix
     - --exit-non-zero-on-fix
     - --preview
+    - --exclude=docs/**/*_.py
 - repo: https://github.com/sphinx-contrib/sphinx-lint
   rev: v1.0.0
   hooks:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 57078fd0..2c7c73c2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -28,6 +28,7 @@ Keep it human-readable, your future self will thank you!
 - Change Changelog CI to run after successful publish
 - pytest for downstream-ci-hpc
 - Update CODEOWNERS
+- Fix pre-commit regex
 
 ### Removed
 

From c5dbb1646a130e7424302c8ea96e38be1d9b3e02 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 20 Oct 2024 19:33:29 +0000
Subject: [PATCH 11/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/anemoi/models/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anemoi/models/__init__.py b/src/anemoi/models/__init__.py
index 7a799a1a..eef2c1d4 100644
--- a/src/anemoi/models/__init__.py
+++ b/src/anemoi/models/__init__.py
@@ -6,4 +6,4 @@
 # nor does it submit to any jurisdiction.
 
 
-from ._version import __version__
+from ._version import __version__ as __version__

From f1057ae1db1dc8107742228622391d01d18ccff6 Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>
Date: Tue, 22 Oct 2024 17:58:09 +0200
Subject: [PATCH 12/18] fix: pin python versions in Github actions (#66)

---
 .github/workflows/python-publish.yml      | 2 +-
 .github/workflows/python-pull-request.yml | 2 +-
 CHANGELOG.md                              | 1 +
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 2f692eda..f5d32232 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -18,7 +18,7 @@ jobs:
     if: ${{ !contains(github.repository, 'private') }}
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.x"]
+        python-version: ["3.9", "3.10", "3.10", "3.11"]
     uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
     with:
       python-version: ${{ matrix.python-version }}
diff --git a/.github/workflows/python-pull-request.yml b/.github/workflows/python-pull-request.yml
index 4b58b126..d4b2379f 100644
--- a/.github/workflows/python-pull-request.yml
+++ b/.github/workflows/python-pull-request.yml
@@ -18,7 +18,7 @@ jobs:
   checks:
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.x"]
+        python-version: ["3.9", "3.10", "3.10", "3.11"]
     uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
     with:
       python-version: ${{ matrix.python-version }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 57078fd0..73a97571 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -28,6 +28,7 @@ Keep it human-readable, your future self will thank you!
 - Change Changelog CI to run after successful publish
 - pytest for downstream-ci-hpc
 - Update CODEOWNERS
+- ci: extened python versions to include 3.11 and 3.12 [#66](https://github.com/ecmwf/anemoi-models/pull/66)
 
 ### Removed
 

From 97c157732796b2053dc8b2ac5f6edf30dd038e7f Mon Sep 17 00:00:00 2001
From: b8raoult <53792887+b8raoult@users.noreply.github.com>
Date: Wed, 23 Oct 2024 16:28:38 +0100
Subject: [PATCH 13/18] Update copyright notice (#67)

---
 CHANGELOG.md                                             | 1 +
 README.md                                                | 2 +-
 src/anemoi/models/commands/hello.py                      | 6 +++---
 src/anemoi/models/data_indices/collection.py             | 5 +++--
 src/anemoi/models/data_indices/index.py                  | 5 +++--
 src/anemoi/models/data_indices/tensor.py                 | 5 +++--
 src/anemoi/models/distributed/graph.py                   | 4 ++--
 src/anemoi/models/distributed/khop_edges.py              | 5 +++--
 src/anemoi/models/distributed/primitives.py              | 5 +++--
 src/anemoi/models/distributed/shapes.py                  | 5 +++--
 src/anemoi/models/distributed/transformer.py             | 5 +++--
 src/anemoi/models/distributed/utils.py                   | 5 +++--
 src/anemoi/models/layers/attention.py                    | 5 +++--
 src/anemoi/models/layers/block.py                        | 5 +++--
 src/anemoi/models/layers/bounding.py                     | 9 +++++++++
 src/anemoi/models/layers/chunk.py                        | 5 +++--
 src/anemoi/models/layers/conv.py                         | 5 +++--
 src/anemoi/models/layers/graph.py                        | 5 +++--
 src/anemoi/models/layers/mapper.py                       | 5 +++--
 src/anemoi/models/layers/mlp.py                          | 5 +++--
 src/anemoi/models/layers/processor.py                    | 5 +++--
 src/anemoi/models/layers/utils.py                        | 5 +++--
 src/anemoi/models/models/encoder_processor_decoder.py    | 5 +++--
 src/anemoi/models/preprocessing/imputer.py               | 5 +++--
 src/anemoi/models/preprocessing/normalizer.py            | 5 +++--
 src/anemoi/models/preprocessing/remapper.py              | 5 +++--
 tests/data_indices/test_collection.py                    | 5 ++++-
 tests/data_indices/test_data_indices.py                  | 5 ++++-
 tests/layers/block/test_block_graphconv.py               | 5 ++++-
 tests/layers/block/test_block_graphtransformer.py        | 5 ++++-
 tests/layers/block/test_block_transformer.py             | 5 ++++-
 tests/layers/chunk/test_chunk_gnn.py                     | 5 ++++-
 tests/layers/chunk/test_chunk_graphtransformer.py        | 5 ++++-
 tests/layers/chunk/test_chunk_transformer.py             | 5 ++++-
 tests/layers/mapper/test_base_mapper.py                  | 5 ++++-
 tests/layers/mapper/test_graphconv_mapper.py             | 5 ++++-
 tests/layers/mapper/test_graphtransformer_mapper.py      | 5 ++++-
 tests/layers/processor/test_base_processor.py            | 5 ++++-
 tests/layers/processor/test_graphconv_processor.py       | 5 ++++-
 .../layers/processor/test_graphtransformer_processor.py  | 5 ++++-
 tests/layers/processor/test_transformer_processor.py     | 5 ++++-
 tests/layers/test_attention.py                           | 5 ++++-
 tests/layers/test_bounding.py                            | 9 +++++++++
 tests/layers/test_graph.py                               | 5 ++++-
 tests/layers/test_mlp.py                                 | 5 ++++-
 tests/models/test_models.py                              | 4 +++-
 tests/preprocessing/test_preprocessor_imputer.py         | 5 +++--
 tests/preprocessing/test_preprocessor_normalizer.py      | 5 +++--
 tests/preprocessing/test_preprocessor_remapper.py        | 5 +++--
 49 files changed, 172 insertions(+), 73 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 73a97571..3292dd7a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ Keep it human-readable, your future self will thank you!
 - pytest for downstream-ci-hpc
 - Update CODEOWNERS
 - ci: extened python versions to include 3.11 and 3.12 [#66](https://github.com/ecmwf/anemoi-models/pull/66)
+- Update copyright notice
 
 ### Removed
 
diff --git a/README.md b/README.md
index 26c7af11..cb578946 100644
--- a/README.md
+++ b/README.md
@@ -22,7 +22,7 @@ $ pip install anemoi-models
 ## License
 
 ```
-Copyright 2022, European Centre for Medium Range Weather Forecasts.
+Copyright 2024, Anemoi contributors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
diff --git a/src/anemoi/models/commands/hello.py b/src/anemoi/models/commands/hello.py
index 12a0495f..d98ecb86 100644
--- a/src/anemoi/models/commands/hello.py
+++ b/src/anemoi/models/commands/hello.py
@@ -1,12 +1,12 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 """Command place holder. Delete when we have real commands.
 
diff --git a/src/anemoi/models/data_indices/collection.py b/src/anemoi/models/data_indices/collection.py
index 266c11a3..2d54544f 100644
--- a/src/anemoi/models/data_indices/collection.py
+++ b/src/anemoi/models/data_indices/collection.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import operator
 
diff --git a/src/anemoi/models/data_indices/index.py b/src/anemoi/models/data_indices/index.py
index 1c8b032c..6cd6507e 100644
--- a/src/anemoi/models/data_indices/index.py
+++ b/src/anemoi/models/data_indices/index.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from anemoi.models.data_indices.tensor import InputTensorIndex
 from anemoi.models.data_indices.tensor import OutputTensorIndex
diff --git a/src/anemoi/models/data_indices/tensor.py b/src/anemoi/models/data_indices/tensor.py
index c7306cf9..f47143de 100644
--- a/src/anemoi/models/data_indices/tensor.py
+++ b/src/anemoi/models/data_indices/tensor.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import torch
 
diff --git a/src/anemoi/models/distributed/graph.py b/src/anemoi/models/distributed/graph.py
index 3d2e3d27..6aab2b96 100644
--- a/src/anemoi/models/distributed/graph.py
+++ b/src/anemoi/models/distributed/graph.py
@@ -1,11 +1,11 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
 
 
 import torch
diff --git a/src/anemoi/models/distributed/khop_edges.py b/src/anemoi/models/distributed/khop_edges.py
index 73397d9a..1d689fdf 100644
--- a/src/anemoi/models/distributed/khop_edges.py
+++ b/src/anemoi/models/distributed/khop_edges.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from typing import Optional
 from typing import Union
diff --git a/src/anemoi/models/distributed/primitives.py b/src/anemoi/models/distributed/primitives.py
index 39b7dea8..7ea3b9ea 100644
--- a/src/anemoi/models/distributed/primitives.py
+++ b/src/anemoi/models/distributed/primitives.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from typing import Optional
 
diff --git a/src/anemoi/models/distributed/shapes.py b/src/anemoi/models/distributed/shapes.py
index 7ba9efb5..01dd8027 100644
--- a/src/anemoi/models/distributed/shapes.py
+++ b/src/anemoi/models/distributed/shapes.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from typing import Optional
 
diff --git a/src/anemoi/models/distributed/transformer.py b/src/anemoi/models/distributed/transformer.py
index 4be22413..78691bba 100644
--- a/src/anemoi/models/distributed/transformer.py
+++ b/src/anemoi/models/distributed/transformer.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from typing import Optional
 
diff --git a/src/anemoi/models/distributed/utils.py b/src/anemoi/models/distributed/utils.py
index e41d145c..b8a6dddf 100644
--- a/src/anemoi/models/distributed/utils.py
+++ b/src/anemoi/models/distributed/utils.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import torch
 from torch import Tensor
diff --git a/src/anemoi/models/layers/attention.py b/src/anemoi/models/layers/attention.py
index 931e098e..d7f54920 100644
--- a/src/anemoi/models/layers/attention.py
+++ b/src/anemoi/models/layers/attention.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 from typing import Optional
diff --git a/src/anemoi/models/layers/block.py b/src/anemoi/models/layers/block.py
index 129e22ed..60446d6c 100644
--- a/src/anemoi/models/layers/block.py
+++ b/src/anemoi/models/layers/block.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 import os
diff --git a/src/anemoi/models/layers/bounding.py b/src/anemoi/models/layers/bounding.py
index 3791ff2c..b168591d 100644
--- a/src/anemoi/models/layers/bounding.py
+++ b/src/anemoi/models/layers/bounding.py
@@ -1,3 +1,12 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
 from __future__ import annotations
 
 from abc import ABC
diff --git a/src/anemoi/models/layers/chunk.py b/src/anemoi/models/layers/chunk.py
index 87d0ac75..5c4fae38 100644
--- a/src/anemoi/models/layers/chunk.py
+++ b/src/anemoi/models/layers/chunk.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 from abc import ABC
diff --git a/src/anemoi/models/layers/conv.py b/src/anemoi/models/layers/conv.py
index e486367f..6b3a767e 100644
--- a/src/anemoi/models/layers/conv.py
+++ b/src/anemoi/models/layers/conv.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from typing import Optional
 
diff --git a/src/anemoi/models/layers/graph.py b/src/anemoi/models/layers/graph.py
index 71703d9f..c7dbefca 100644
--- a/src/anemoi/models/layers/graph.py
+++ b/src/anemoi/models/layers/graph.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import einops
 import torch
diff --git a/src/anemoi/models/layers/mapper.py b/src/anemoi/models/layers/mapper.py
index dc5c1f08..1ae45031 100644
--- a/src/anemoi/models/layers/mapper.py
+++ b/src/anemoi/models/layers/mapper.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 from abc import ABC
diff --git a/src/anemoi/models/layers/mlp.py b/src/anemoi/models/layers/mlp.py
index 4de53001..4a1e7957 100644
--- a/src/anemoi/models/layers/mlp.py
+++ b/src/anemoi/models/layers/mlp.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 
diff --git a/src/anemoi/models/layers/processor.py b/src/anemoi/models/layers/processor.py
index f35845ae..4fd32311 100644
--- a/src/anemoi/models/layers/processor.py
+++ b/src/anemoi/models/layers/processor.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from abc import ABC
 from typing import Optional
diff --git a/src/anemoi/models/layers/utils.py b/src/anemoi/models/layers/utils.py
index 90cdcc9a..e243874a 100644
--- a/src/anemoi/models/layers/utils.py
+++ b/src/anemoi/models/layers/utils.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 from torch import Tensor
 from torch import nn
diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py
index c77db6e8..bdb6260e 100644
--- a/src/anemoi/models/models/encoder_processor_decoder.py
+++ b/src/anemoi/models/models/encoder_processor_decoder.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 from typing import Optional
diff --git a/src/anemoi/models/preprocessing/imputer.py b/src/anemoi/models/preprocessing/imputer.py
index 6ef5adbe..5f9c1b9d 100644
--- a/src/anemoi/models/preprocessing/imputer.py
+++ b/src/anemoi/models/preprocessing/imputer.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 from abc import ABC
diff --git a/src/anemoi/models/preprocessing/normalizer.py b/src/anemoi/models/preprocessing/normalizer.py
index ee6a4f5e..aacc7164 100644
--- a/src/anemoi/models/preprocessing/normalizer.py
+++ b/src/anemoi/models/preprocessing/normalizer.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 import warnings
diff --git a/src/anemoi/models/preprocessing/remapper.py b/src/anemoi/models/preprocessing/remapper.py
index a79e2af8..cc888222 100644
--- a/src/anemoi/models/preprocessing/remapper.py
+++ b/src/anemoi/models/preprocessing/remapper.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import logging
 from abc import ABC
diff --git a/tests/data_indices/test_collection.py b/tests/data_indices/test_collection.py
index 8505f8aa..9e2e4b06 100644
--- a/tests/data_indices/test_collection.py
+++ b/tests/data_indices/test_collection.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from omegaconf import DictConfig
diff --git a/tests/data_indices/test_data_indices.py b/tests/data_indices/test_data_indices.py
index e447832f..44cf9b46 100644
--- a/tests/data_indices/test_data_indices.py
+++ b/tests/data_indices/test_data_indices.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 
diff --git a/tests/layers/block/test_block_graphconv.py b/tests/layers/block/test_block_graphconv.py
index 34083e6e..fe89cb63 100644
--- a/tests/layers/block/test_block_graphconv.py
+++ b/tests/layers/block/test_block_graphconv.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 from hypothesis import given
 from hypothesis import settings
 from hypothesis import strategies as st
diff --git a/tests/layers/block/test_block_graphtransformer.py b/tests/layers/block/test_block_graphtransformer.py
index 3d1c1da4..a6162046 100644
--- a/tests/layers/block/test_block_graphtransformer.py
+++ b/tests/layers/block/test_block_graphtransformer.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import importlib
 
 import pytest
diff --git a/tests/layers/block/test_block_transformer.py b/tests/layers/block/test_block_transformer.py
index 2e63386c..46541e08 100644
--- a/tests/layers/block/test_block_transformer.py
+++ b/tests/layers/block/test_block_transformer.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import logging
 
 import torch
diff --git a/tests/layers/chunk/test_chunk_gnn.py b/tests/layers/chunk/test_chunk_gnn.py
index 64cc04b7..f8ac8a12 100644
--- a/tests/layers/chunk/test_chunk_gnn.py
+++ b/tests/layers/chunk/test_chunk_gnn.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 
 from anemoi.models.layers.block import GraphConvProcessorBlock
diff --git a/tests/layers/chunk/test_chunk_graphtransformer.py b/tests/layers/chunk/test_chunk_graphtransformer.py
index 1249e237..2f93cc91 100644
--- a/tests/layers/chunk/test_chunk_graphtransformer.py
+++ b/tests/layers/chunk/test_chunk_graphtransformer.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 
 from anemoi.models.layers.block import GraphTransformerProcessorBlock
diff --git a/tests/layers/chunk/test_chunk_transformer.py b/tests/layers/chunk/test_chunk_transformer.py
index 5449e97c..86989486 100644
--- a/tests/layers/chunk/test_chunk_transformer.py
+++ b/tests/layers/chunk/test_chunk_transformer.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 
 from anemoi.models.layers.block import TransformerProcessorBlock
diff --git a/tests/layers/mapper/test_base_mapper.py b/tests/layers/mapper/test_base_mapper.py
index 3cc4ef0f..85f4b644 100644
--- a/tests/layers/mapper/test_base_mapper.py
+++ b/tests/layers/mapper/test_base_mapper.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from torch_geometric.data import HeteroData
diff --git a/tests/layers/mapper/test_graphconv_mapper.py b/tests/layers/mapper/test_graphconv_mapper.py
index 4be7130c..1a756989 100644
--- a/tests/layers/mapper/test_graphconv_mapper.py
+++ b/tests/layers/mapper/test_graphconv_mapper.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from torch import nn
diff --git a/tests/layers/mapper/test_graphtransformer_mapper.py b/tests/layers/mapper/test_graphtransformer_mapper.py
index c872422d..dece0e22 100644
--- a/tests/layers/mapper/test_graphtransformer_mapper.py
+++ b/tests/layers/mapper/test_graphtransformer_mapper.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from torch import nn
diff --git a/tests/layers/processor/test_base_processor.py b/tests/layers/processor/test_base_processor.py
index 4af3c7bc..bd423d1c 100644
--- a/tests/layers/processor/test_base_processor.py
+++ b/tests/layers/processor/test_base_processor.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 
 from anemoi.models.layers.processor import BaseProcessor
diff --git a/tests/layers/processor/test_graphconv_processor.py b/tests/layers/processor/test_graphconv_processor.py
index 2505515c..e847d64e 100644
--- a/tests/layers/processor/test_graphconv_processor.py
+++ b/tests/layers/processor/test_graphconv_processor.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from torch_geometric.data import HeteroData
diff --git a/tests/layers/processor/test_graphtransformer_processor.py b/tests/layers/processor/test_graphtransformer_processor.py
index dfba417e..95ba1c45 100644
--- a/tests/layers/processor/test_graphtransformer_processor.py
+++ b/tests/layers/processor/test_graphtransformer_processor.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from torch_geometric.data import HeteroData
diff --git a/tests/layers/processor/test_transformer_processor.py b/tests/layers/processor/test_transformer_processor.py
index 305af413..b94ff63f 100644
--- a/tests/layers/processor/test_transformer_processor.py
+++ b/tests/layers/processor/test_transformer_processor.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 
diff --git a/tests/layers/test_attention.py b/tests/layers/test_attention.py
index 9457317d..a1b40540 100644
--- a/tests/layers/test_attention.py
+++ b/tests/layers/test_attention.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import hypothesis.strategies as st
 import pytest
 import torch
diff --git a/tests/layers/test_bounding.py b/tests/layers/test_bounding.py
index 87619cdc..eddbae2d 100644
--- a/tests/layers/test_bounding.py
+++ b/tests/layers/test_bounding.py
@@ -1,3 +1,12 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
 import pytest
 import torch
 from anemoi.utils.config import DotDict
diff --git a/tests/layers/test_graph.py b/tests/layers/test_graph.py
index a17ebfd4..58674bd5 100644
--- a/tests/layers/test_graph.py
+++ b/tests/layers/test_graph.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 from torch import nn
diff --git a/tests/layers/test_mlp.py b/tests/layers/test_mlp.py
index e47a0b98..a5e93892 100644
--- a/tests/layers/test_mlp.py
+++ b/tests/layers/test_mlp.py
@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import pytest
 import torch
 
diff --git a/tests/models/test_models.py b/tests/models/test_models.py
index 3b4eb3f4..f005f3ff 100644
--- a/tests/models/test_models.py
+++ b/tests/models/test_models.py
@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
diff --git a/tests/preprocessing/test_preprocessor_imputer.py b/tests/preprocessing/test_preprocessor_imputer.py
index 5218efcb..a22e261c 100644
--- a/tests/preprocessing/test_preprocessor_imputer.py
+++ b/tests/preprocessing/test_preprocessor_imputer.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import numpy as np
 import pytest
diff --git a/tests/preprocessing/test_preprocessor_normalizer.py b/tests/preprocessing/test_preprocessor_normalizer.py
index 3a2327ea..ad407413 100644
--- a/tests/preprocessing/test_preprocessor_normalizer.py
+++ b/tests/preprocessing/test_preprocessor_normalizer.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import numpy as np
 import pytest
diff --git a/tests/preprocessing/test_preprocessor_remapper.py b/tests/preprocessing/test_preprocessor_remapper.py
index 86bdfde1..a0ece2a3 100644
--- a/tests/preprocessing/test_preprocessor_remapper.py
+++ b/tests/preprocessing/test_preprocessor_remapper.py
@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
+
 
 import pytest
 import torch

From 9dc1d7950b598689f8b148692c5088851f16ac2d Mon Sep 17 00:00:00 2001
From: Harrison Cook <harrison.cook@ecmwf.int>
Date: Fri, 25 Oct 2024 10:05:04 +0000
Subject: [PATCH 14/18]  Fix __version__ import  - Use standard try-except

---
 CHANGELOG.md                  | 1 +
 src/anemoi/models/__init__.py | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3292dd7a..5750c66b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,7 @@ Keep it human-readable, your future self will thank you!
 - Update CODEOWNERS
 - ci: extened python versions to include 3.11 and 3.12 [#66](https://github.com/ecmwf/anemoi-models/pull/66)
 - Update copyright notice
+- Fix `__version__` import in init
 
 ### Removed
 
diff --git a/src/anemoi/models/__init__.py b/src/anemoi/models/__init__.py
index 7a799a1a..9733be26 100644
--- a/src/anemoi/models/__init__.py
+++ b/src/anemoi/models/__init__.py
@@ -6,4 +6,10 @@
 # nor does it submit to any jurisdiction.
 
 
-from ._version import __version__
+try:
+    # NOTE: the `_version.py` file must not be present in the git repository
+    #   as it is generated by setuptools at install time
+    from ._version import __version__  # type: ignore
+except ImportError:  # pragma: no cover
+    # Local copy or not installed with setuptools
+    __version__ = "999"

From deba13a3938e2d37c42b82f4c2e42bd2c17c5b17 Mon Sep 17 00:00:00 2001
From: Matthew Chantry <chantry.m.j@gmail.com>
Date: Mon, 28 Oct 2024 11:11:16 +0000
Subject: [PATCH 15/18] Add contributors (#69)

* Add contributors

* fix: update changelog

---------

Co-authored-by: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
---
 CHANGELOG.md    |  1 +
 CONTRIBUTORS.md | 13 +++++++++++++
 2 files changed, 14 insertions(+)
 create mode 100644 CONTRIBUTORS.md

diff --git a/CHANGELOG.md b/CHANGELOG.md
index acf1aaed..f68f20d4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,7 @@ Keep it human-readable, your future self will thank you!
 - configurabilty of the dropout probability in the the MultiHeadSelfAttention module
 - Variable Bounding as configurable model layers [#13](https://github.com/ecmwf/anemoi-models/issues/13)
 - GraphTransformerMapperBlock chunking to reduce memory usage during inference [#46](https://github.com/ecmwf/anemoi-models/pull/46)
+- Contributors file [#69](https://github.com/ecmwf/anemoi-models/pull/69)
 
 ### Changed
 - Bugfixes for CI
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
new file mode 100644
index 00000000..5a6d0c13
--- /dev/null
+++ b/CONTRIBUTORS.md
@@ -0,0 +1,13 @@
+## How to Contribute
+
+Please see the [read the docs](https://anemoi-training.readthedocs.io/en/latest/dev/contributing.html).
+
+
+## Contributors
+
+Thank you to all the wonderful people who have contributed to Anemoi. Contributions can come in many forms, including code, documentation, bug reports, feature suggestions, design, and more. A list of code-based contributors can be found [here](https://github.com/ecmwf/anemoi-models/graphs/contributors).
+
+
+## Contributing Organisations
+
+Significant contributions have been made by the following organisations: [DWD](https://www.dwd.de/), [MET Norway](https://www.met.no/), [MeteoSwiss](https://www.meteoswiss.admin.ch/), [RMI](https://www.meteo.be/) & [ECMWF](https://www.ecmwf.int/)

From 52c232bfc14229937a8c1c722d177a7226edc1ac Mon Sep 17 00:00:00 2001
From: Jesper Dramsch <jesper.dramsch@ecmwf.int>
Date: Mon, 11 Nov 2024 09:22:09 +0100
Subject: [PATCH 16/18] Update .pre-commit-config.yaml

---
 .pre-commit-config.yaml | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 4de59323..bfaf353f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -59,11 +59,6 @@ repos:
   hooks:
   - id: rstfmt
     exclude: 'cli/.*' # Because we use argparse
-- repo: https://github.com/b8raoult/pre-commit-docconvert
-  rev: "0.1.5"
-  hooks:
-  - id: docconvert
-    args: ["numpy"]
 - repo: https://github.com/tox-dev/pyproject-fmt
   rev: "2.2.4"
   hooks:

From 8874571592529e1e4d73278503459151f37207d0 Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>
Date: Mon, 11 Nov 2024 16:56:22 +0100
Subject: [PATCH 17/18] Refactor node attributes (#64)

* feat: add NamedNodeAttributes

* feat: use NamedNodesAttributes in AnemoiModelEncProcDec

* fix: update changelog

* fix: add tests

* feat: drop unused attrs + type hints

* fix: homogeneise

* fix: update docstring

* fix: more docstrings
---
 CHANGELOG.md                                  |  1 +
 src/anemoi/models/layers/graph.py             | 72 ++++++++++++++++-
 .../models/encoder_processor_decoder.py       | 65 ++++-----------
 tests/layers/test_graph.py                    | 79 +++++++++++++++++++
 4 files changed, 165 insertions(+), 52 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f68f20d4..470cfaf8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,7 @@ Keep it human-readable, your future self will thank you!
 - configurabilty of the dropout probability in the the MultiHeadSelfAttention module
 - Variable Bounding as configurable model layers [#13](https://github.com/ecmwf/anemoi-models/issues/13)
 - GraphTransformerMapperBlock chunking to reduce memory usage during inference [#46](https://github.com/ecmwf/anemoi-models/pull/46)
+- New `NamedNodesAttributes` class to handle node attributes in a more flexible way [#64](https://github.com/ecmwf/anemoi-models/pull/64)
 - Contributors file [#69](https://github.com/ecmwf/anemoi-models/pull/69)
 
 ### Changed
diff --git a/src/anemoi/models/layers/graph.py b/src/anemoi/models/layers/graph.py
index c7dbefca..a934d32b 100644
--- a/src/anemoi/models/layers/graph.py
+++ b/src/anemoi/models/layers/graph.py
@@ -12,6 +12,7 @@
 import torch
 from torch import Tensor
 from torch import nn
+from torch_geometric.data import HeteroData
 
 
 class TrainableTensor(nn.Module):
@@ -36,8 +37,77 @@ def __init__(self, tensor_size: int, trainable_size: int) -> None:
     def forward(self, x: Tensor, batch_size: int) -> Tensor:
         latent = [einops.repeat(x, "e f -> (repeat e) f", repeat=batch_size)]
         if self.trainable is not None:
-            latent.append(einops.repeat(self.trainable, "e f -> (repeat e) f", repeat=batch_size))
+            latent.append(einops.repeat(self.trainable.to(x.device), "e f -> (repeat e) f", repeat=batch_size))
         return torch.cat(
             latent,
             dim=-1,  # feature dimension
         )
+
+
+class NamedNodesAttributes(nn.Module):
+    """Named Nodes Attributes information.
+
+    Attributes
+    ----------
+    num_nodes : dict[str, int]
+        Number of nodes for each group of nodes.
+    attr_ndims : dict[str, int]
+        Total dimension of node attributes (non-trainable + trainable) for each group of nodes.
+    trainable_tensors : nn.ModuleDict
+        Dictionary of trainable tensors for each group of nodes.
+
+    Methods
+    -------
+    get_coordinates(self, name: str) -> Tensor
+        Get the coordinates of a set of nodes.
+    forward( self, name: str, batch_size: int) -> Tensor
+        Get the node attributes to be passed trough the graph neural network.
+    """
+
+    num_nodes: dict[str, int]
+    attr_ndims: dict[str, int]
+    trainable_tensors: dict[str, TrainableTensor]
+
+    def __init__(self, num_trainable_params: int, graph_data: HeteroData) -> None:
+        """Initialize NamedNodesAttributes."""
+        super().__init__()
+
+        self.define_fixed_attributes(graph_data, num_trainable_params)
+
+        self.trainable_tensors = nn.ModuleDict()
+        for nodes_name, nodes in graph_data.node_items():
+            self.register_coordinates(nodes_name, nodes.x)
+            self.register_tensor(nodes_name, num_trainable_params)
+
+    def define_fixed_attributes(self, graph_data: HeteroData, num_trainable_params: int) -> None:
+        """Define fixed attributes."""
+        nodes_names = list(graph_data.node_types)
+        self.num_nodes = {nodes_name: graph_data[nodes_name].num_nodes for nodes_name in nodes_names}
+        self.attr_ndims = {
+            nodes_name: 2 * graph_data[nodes_name].x.shape[1] + num_trainable_params for nodes_name in nodes_names
+        }
+
+    def register_coordinates(self, name: str, node_coords: Tensor) -> None:
+        """Register coordinates."""
+        sin_cos_coords = torch.cat([torch.sin(node_coords), torch.cos(node_coords)], dim=-1)
+        self.register_buffer(f"latlons_{name}", sin_cos_coords, persistent=True)
+
+    def get_coordinates(self, name: str) -> Tensor:
+        """Return original coordinates."""
+        sin_cos_coords = getattr(self, f"latlons_{name}")
+        ndim = sin_cos_coords.shape[1] // 2
+        sin_values = sin_cos_coords[:, :ndim]
+        cos_values = sin_cos_coords[:, ndim:]
+        return torch.atan2(sin_values, cos_values)
+
+    def register_tensor(self, name: str, num_trainable_params: int) -> None:
+        """Register a trainable tensor."""
+        self.trainable_tensors[name] = TrainableTensor(self.num_nodes[name], num_trainable_params)
+
+    def forward(self, name: str, batch_size: int) -> Tensor:
+        """Returns the node attributes to be passed trough the graph neural network.
+
+        It includes both the coordinates and the trainable parameters.
+        """
+        latlons = getattr(self, f"latlons_{name}")
+        return self.trainable_tensors[name](latlons, batch_size)
diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py
index bdb6260e..c67c8c03 100644
--- a/src/anemoi/models/models/encoder_processor_decoder.py
+++ b/src/anemoi/models/models/encoder_processor_decoder.py
@@ -22,7 +22,7 @@
 from torch_geometric.data import HeteroData
 
 from anemoi.models.distributed.shapes import get_shape_shards
-from anemoi.models.layers.graph import TrainableTensor
+from anemoi.models.layers.graph import NamedNodesAttributes
 
 LOGGER = logging.getLogger(__name__)
 
@@ -56,33 +56,24 @@ def __init__(
 
         self._calculate_shapes_and_indices(data_indices)
         self._assert_matching_indices(data_indices)
-
-        self.multi_step = model_config.training.multistep_input
-
-        self._define_tensor_sizes(model_config)
-
-        # Create trainable tensors
-        self._create_trainable_attributes()
-
-        # Register lat/lon of nodes
-        self._register_latlon("data", self._graph_name_data)
-        self._register_latlon("hidden", self._graph_name_hidden)
-
         self.data_indices = data_indices
 
+        self.multi_step = model_config.training.multistep_input
         self.num_channels = model_config.model.num_channels
 
-        input_dim = self.multi_step * self.num_input_channels + self.latlons_data.shape[1] + self.trainable_data_size
+        self.node_attributes = NamedNodesAttributes(model_config.model.trainable_parameters.hidden, self._graph_data)
+
+        input_dim = self.multi_step * self.num_input_channels + self.node_attributes.attr_ndims[self._graph_name_data]
 
         # Encoder data -> hidden
         self.encoder = instantiate(
             model_config.model.encoder,
             in_channels_src=input_dim,
-            in_channels_dst=self.latlons_hidden.shape[1] + self.trainable_hidden_size,
+            in_channels_dst=self.node_attributes.attr_ndims[self._graph_name_hidden],
             hidden_dim=self.num_channels,
             sub_graph=self._graph_data[(self._graph_name_data, "to", self._graph_name_hidden)],
-            src_grid_size=self._data_grid_size,
-            dst_grid_size=self._hidden_grid_size,
+            src_grid_size=self.node_attributes.num_nodes[self._graph_name_data],
+            dst_grid_size=self.node_attributes.num_nodes[self._graph_name_hidden],
         )
 
         # Processor hidden -> hidden
@@ -90,8 +81,8 @@ def __init__(
             model_config.model.processor,
             num_channels=self.num_channels,
             sub_graph=self._graph_data[(self._graph_name_hidden, "to", self._graph_name_hidden)],
-            src_grid_size=self._hidden_grid_size,
-            dst_grid_size=self._hidden_grid_size,
+            src_grid_size=self.node_attributes.num_nodes[self._graph_name_hidden],
+            dst_grid_size=self.node_attributes.num_nodes[self._graph_name_hidden],
         )
 
         # Decoder hidden -> data
@@ -102,8 +93,8 @@ def __init__(
             hidden_dim=self.num_channels,
             out_channels_dst=self.num_output_channels,
             sub_graph=self._graph_data[(self._graph_name_hidden, "to", self._graph_name_data)],
-            src_grid_size=self._hidden_grid_size,
-            dst_grid_size=self._data_grid_size,
+            src_grid_size=self.node_attributes.num_nodes[self._graph_name_hidden],
+            dst_grid_size=self.node_attributes.num_nodes[self._graph_name_data],
         )
 
         # Instantiation of model output bounding functions (e.g., to ensure outputs like TP are positive definite)
@@ -133,34 +124,6 @@ def _assert_matching_indices(self, data_indices: dict) -> None:
             self._internal_output_idx,
         ), f"Internal model indices must match {self._internal_input_idx} != {self._internal_output_idx}"
 
-    def _define_tensor_sizes(self, config: DotDict) -> None:
-        self._data_grid_size = self._graph_data[self._graph_name_data].num_nodes
-        self._hidden_grid_size = self._graph_data[self._graph_name_hidden].num_nodes
-
-        self.trainable_data_size = config.model.trainable_parameters.data
-        self.trainable_hidden_size = config.model.trainable_parameters.hidden
-
-    def _register_latlon(self, name: str, nodes: str) -> None:
-        """Register lat/lon buffers.
-
-        Parameters
-        ----------
-        name : str
-            Name to store the lat-lon coordinates of the nodes.
-        nodes : str
-            Name of nodes to map
-        """
-        coords = self._graph_data[nodes].x
-        sin_cos_coords = torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
-        self.register_buffer(f"latlons_{name}", sin_cos_coords, persistent=True)
-
-    def _create_trainable_attributes(self) -> None:
-        """Create all trainable attributes."""
-        self.trainable_data = TrainableTensor(trainable_size=self.trainable_data_size, tensor_size=self._data_grid_size)
-        self.trainable_hidden = TrainableTensor(
-            trainable_size=self.trainable_hidden_size, tensor_size=self._hidden_grid_size
-        )
-
     def _run_mapper(
         self,
         mapper: nn.Module,
@@ -210,12 +173,12 @@ def forward(self, x: Tensor, model_comm_group: Optional[ProcessGroup] = None) ->
         x_data_latent = torch.cat(
             (
                 einops.rearrange(x, "batch time ensemble grid vars -> (batch ensemble grid) (time vars)"),
-                self.trainable_data(self.latlons_data, batch_size=batch_size),
+                self.node_attributes(self._graph_name_data, batch_size=batch_size),
             ),
             dim=-1,  # feature dimension
         )
 
-        x_hidden_latent = self.trainable_hidden(self.latlons_hidden, batch_size=batch_size)
+        x_hidden_latent = self.node_attributes(self._graph_name_hidden, batch_size=batch_size)
 
         # get shard shapes
         shard_shapes_data = get_shape_shards(x_data_latent, 0, model_comm_group)
diff --git a/tests/layers/test_graph.py b/tests/layers/test_graph.py
index 58674bd5..66456d6c 100644
--- a/tests/layers/test_graph.py
+++ b/tests/layers/test_graph.py
@@ -8,10 +8,14 @@
 # nor does it submit to any jurisdiction.
 
 
+import einops
+import numpy as np
 import pytest
 import torch
 from torch import nn
+from torch_geometric.data import HeteroData
 
+from anemoi.models.layers.graph import NamedNodesAttributes
 from anemoi.models.layers.graph import TrainableTensor
 
 
@@ -62,3 +66,78 @@ def test_forward_no_trainable(self, init, x):
         batch_size = 5
         output = trainable_tensor(x, batch_size)
         assert output.shape == (batch_size * x.shape[0], tensor_size + trainable_size)
+
+
+class TestNamedNodesAttributes:
+    """Test suite for the NamedNodesAttributes class.
+
+    This class contains test cases to verify the functionality of the NamedNodesAttributes class,
+    including initialization, attribute registration, and forward pass operations.
+    """
+
+    nodes_names: list[str] = ["nodes1", "nodes2"]
+    ndim: int = 2
+    num_trainable_params: int = 8
+
+    @pytest.fixture
+    def graph_data(self):
+        graph = HeteroData()
+        for i, nodes_name in enumerate(TestNamedNodesAttributes.nodes_names):
+            graph[nodes_name].x = TestNamedNodesAttributes.get_n_random_coords(10 + 5 ** (i + 1))
+        return graph
+
+    @staticmethod
+    def get_n_random_coords(n: int) -> torch.Tensor:
+        coords = torch.rand(n, TestNamedNodesAttributes.ndim)
+        coords[:, 0] = np.pi * (coords[:, 0] - 1 / 2)
+        coords[:, 1] = 2 * np.pi * coords[:, 1]
+        return coords
+
+    @pytest.fixture
+    def nodes_attributes(self, graph_data: HeteroData) -> NamedNodesAttributes:
+        return NamedNodesAttributes(TestNamedNodesAttributes.num_trainable_params, graph_data)
+
+    def test_init(self, nodes_attributes):
+        assert isinstance(nodes_attributes, NamedNodesAttributes)
+
+        for nodes_name in self.nodes_names:
+            assert isinstance(nodes_attributes.num_nodes[nodes_name], int)
+            assert (
+                nodes_attributes.attr_ndims[nodes_name] - 2 * TestNamedNodesAttributes.ndim
+                == TestNamedNodesAttributes.num_trainable_params
+            )
+            assert isinstance(nodes_attributes.trainable_tensors[nodes_name], TrainableTensor)
+
+    def test_forward(self, nodes_attributes, graph_data):
+        batch_size = 3
+        for nodes_name in self.nodes_names:
+            output = nodes_attributes(nodes_name, batch_size)
+
+            expected_shape = (
+                batch_size * graph_data[nodes_name].num_nodes,
+                2 * TestNamedNodesAttributes.ndim + TestNamedNodesAttributes.num_trainable_params,
+            )
+            assert output.shape == expected_shape
+
+            # Check if the first part of the output matches the sin-cos transformed coordinates
+            latlons = getattr(nodes_attributes, f"latlons_{nodes_name}")
+            repeated_latlons = einops.repeat(latlons, "n f -> (b n) f", b=batch_size)
+            assert torch.allclose(output[:, : 2 * TestNamedNodesAttributes.ndim], repeated_latlons)
+
+            # Check if the last part of the output is trainable (requires grad)
+            assert output[:, 2 * TestNamedNodesAttributes.ndim :].requires_grad
+
+    def test_forward_no_trainable(self, graph_data):
+        no_trainable_attributes = NamedNodesAttributes(0, graph_data)
+        batch_size = 2
+
+        for nodes_name in self.nodes_names:
+            output = no_trainable_attributes(nodes_name, batch_size)
+
+            expected_shape = batch_size * graph_data[nodes_name].num_nodes, 2 * TestNamedNodesAttributes.ndim
+            assert output.shape == expected_shape
+
+            # Check if the output exactly matches the sin-cos transformed coordinates
+            latlons = getattr(no_trainable_attributes, f"latlons_{nodes_name}")
+            repeated_latlons = einops.repeat(latlons, "n f -> (b n) f", b=batch_size)
+            assert torch.allclose(output, repeated_latlons)

From 0e03d33167fe74e99463f0574557f5a1c32b82cf Mon Sep 17 00:00:00 2001
From: b8raoult <53792887+b8raoult@users.noreply.github.com>
Date: Tue, 12 Nov 2024 16:47:51 +0000
Subject: [PATCH 18/18] missing copyrigths (#71)

* Address missing copyrights

---------

Co-authored-by: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>
---
 CHANGELOG.md                                | 1 +
 docs/conf.py                                | 4 ++--
 pyproject.toml                              | 5 ++---
 src/anemoi/models/__init__.py               | 4 +++-
 src/anemoi/models/__main__.py               | 5 ++---
 src/anemoi/models/commands/__init__.py      | 5 ++---
 src/anemoi/models/data_indices/__init__.py  | 8 ++++++++
 src/anemoi/models/distributed/__init__.py   | 8 ++++++++
 src/anemoi/models/interface/__init__.py     | 4 ++--
 src/anemoi/models/layers/__init__.py        | 8 ++++++++
 src/anemoi/models/models/__init__.py        | 8 ++++++++
 src/anemoi/models/preprocessing/__init__.py | 4 ++--
 12 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 470cfaf8..07ee5709 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,6 +34,7 @@ Keep it human-readable, your future self will thank you!
 - ci: extened python versions to include 3.11 and 3.12 [#66](https://github.com/ecmwf/anemoi-models/pull/66)
 - Update copyright notice
 - Fix `__version__` import in init
+- Fix missing copyrights [#71](https://github.com/ecmwf/anemoi-models/pull/71)
 
 ### Removed
 
diff --git a/docs/conf.py b/docs/conf.py
index 8a5d5688..32e7531c 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -29,7 +29,7 @@
 
 project = "Anemoi Models"
 
-author = "ECMWF"
+author = "Anemoi contributors"
 
 year = datetime.datetime.now().year
 if year == 2024:
@@ -37,7 +37,7 @@
 else:
     years = "2024-%s" % (year,)
 
-copyright = "%s, ECMWF" % (years,)
+copyright = "%s, Anemoi contributors" % (years,)
 
 try:
     from anemoi.models._version import __version__
diff --git a/pyproject.toml b/pyproject.toml
index 214f82c1..ba0b9d33 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,13 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-# https://packaging.python.org/en/latest/guides/writing-pyproject-toml/
-
 [build-system]
 build-backend = "setuptools.build_meta"
 
diff --git a/src/anemoi/models/__init__.py b/src/anemoi/models/__init__.py
index bdd630a2..019edd80 100644
--- a/src/anemoi/models/__init__.py
+++ b/src/anemoi/models/__init__.py
@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
diff --git a/src/anemoi/models/__main__.py b/src/anemoi/models/__main__.py
index be940c27..0057b940 100644
--- a/src/anemoi/models/__main__.py
+++ b/src/anemoi/models/__main__.py
@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
 
 from anemoi.utils.cli import cli_main
 from anemoi.utils.cli import make_parser
diff --git a/src/anemoi/models/commands/__init__.py b/src/anemoi/models/commands/__init__.py
index cebb5395..e5e2219d 100644
--- a/src/anemoi/models/commands/__init__.py
+++ b/src/anemoi/models/commands/__init__.py
@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
 
 import os
 
diff --git a/src/anemoi/models/data_indices/__init__.py b/src/anemoi/models/data_indices/__init__.py
index e69de29b..c167afa2 100644
--- a/src/anemoi/models/data_indices/__init__.py
+++ b/src/anemoi/models/data_indices/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
diff --git a/src/anemoi/models/distributed/__init__.py b/src/anemoi/models/distributed/__init__.py
index e69de29b..c167afa2 100644
--- a/src/anemoi/models/distributed/__init__.py
+++ b/src/anemoi/models/distributed/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py
index aba62a23..25b7852a 100644
--- a/src/anemoi/models/interface/__init__.py
+++ b/src/anemoi/models/interface/__init__.py
@@ -1,11 +1,11 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
 
 import uuid
 
diff --git a/src/anemoi/models/layers/__init__.py b/src/anemoi/models/layers/__init__.py
index e69de29b..c167afa2 100644
--- a/src/anemoi/models/layers/__init__.py
+++ b/src/anemoi/models/layers/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
diff --git a/src/anemoi/models/models/__init__.py b/src/anemoi/models/models/__init__.py
index e69de29b..c167afa2 100644
--- a/src/anemoi/models/models/__init__.py
+++ b/src/anemoi/models/models/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
diff --git a/src/anemoi/models/preprocessing/__init__.py b/src/anemoi/models/preprocessing/__init__.py
index cc2cb4f8..4c7f6153 100644
--- a/src/anemoi/models/preprocessing/__init__.py
+++ b/src/anemoi/models/preprocessing/__init__.py
@@ -1,11 +1,11 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-#
 
 import logging
 from typing import Optional