Add mart.nn.Get() to extract a value from the kwargs dict.

mart/nn/nn.py

@@ -13,7 +13,7 @@
 
 import torch
 
-__all__ = ["GroupNorm32", "SequentialDict", "ReturnKwargs", "CallWith", "Sum"]
+__all__ = ["GroupNorm32", "SequentialDict", "ReturnKwargs", "CallWith", "Sum", "Get"]
 
 logger = logging.getLogger(__name__)
 
@@ -300,3 +300,19 @@ def __init__(self):
 
     def forward(self, *args):
         return sum(args)
+
+
+class Get:
+    """Get a value from the kwargs dictionary by key.
+
+    The key can be a path to a nested dictionary, concatenated by dots. For example,
+    `Get(key="a.b")(a={"b": 1}) == 1`.
+    """
+
+    def __init__(self, key):
+        self.key = key
+
+    def __call__(self, **kwargs):
+        # Add support to nested dicts.
+        kwargs = DotDict(kwargs)
+        return kwargs[self.key]