add and adapt type hints

Author: vpratz

bayesflow/adapters/transforms/convert_dtype.py

@@ -32,8 +32,8 @@ def get_config(self) -> dict:
         }
         return serialize(config)
 
-    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+    def forward(self, data: np.ndarray | dict, **kwargs) -> np.ndarray | dict:
         return map_structure(lambda d: d.astype(self.to_dtype, copy=False), data)
 
-    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+    def inverse(self, data: np.ndarray | dict, **kwargs) -> np.ndarray | dict:
         return map_structure(lambda d: d.astype(self.from_dtype, copy=False), data)

bayesflow/utils/tree.py

@@ -1,4 +1,5 @@
 import optree
+from typing import Callable
 
 
 def flatten_shape(structure):
@@ -14,7 +15,7 @@ def is_shape_tuple(x):
     return leaves
 
 
-def map_dict(func, dictionary):
+def map_dict(func: Callable, dictionary: dict) -> dict:
     """Applies a function to all leaves of a (possibly nested) dictionary.
 
     Parameters