From 7c0daa7a9990004d03a830af4a40363ac7514297 Mon Sep 17 00:00:00 2001
From: japols <jan.polster@ecmwf.int>
Date: Tue, 1 Oct 2024 09:48:40 +0000
Subject: [PATCH] docs: add chunking documentation, rename
 ANEMOI_INFERENCE_NUM_CHUNKS

---
 CHANGELOG.md                      |  5 ++---
 docs/modules/layers.rst           | 14 ++++++++++++++
 src/anemoi/models/layers/block.py |  4 +++-
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 42393f9..c660ac8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,8 @@ Keep it human-readable, your future self will thank you!
 ## [Unreleased](https://github.com/ecmwf/anemoi-models/compare/0.3.0...HEAD)
 
 ### Added
+
+- GraphTransformerMapperBlock chunking to reduce memory usage during inference [#46](https://github.com/ecmwf/anemoi-models/pull/46)
 - Codeowners file
 - Pygrep precommit hooks
 - Docsig precommit hooks
@@ -18,9 +20,6 @@ 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)
 
-- CI workflow to update the changelog on release
-- Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords.
-
 ### Changed
 - Bugfixes for CI
 
diff --git a/docs/modules/layers.rst b/docs/modules/layers.rst
index ff21b5a..8082e47 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/layers/block.py b/src/anemoi/models/layers/block.py
index 6083a82..129e22e 100644
--- a/src/anemoi/models/layers/block.py
+++ b/src/anemoi/models/layers/block.py
@@ -33,7 +33,9 @@
 from anemoi.models.layers.mlp import MLP
 
 LOGGER = logging.getLogger(__name__)
-NUM_CHUNKS_INFERENCE = int(os.environ.get("ANEMOI_NUM_CHUNKS_INFERENCE", "1"))
+
+# 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):