docs: add chunking documentation, rename ANEMOI_INFERENCE_NUM_CHUNKS

CHANGELOG.md

@@ -11,16 +11,15 @@ 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
 - 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)
 
-- 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
 

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
 *********

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):