Experiment 01 update (#236)

* mod_config bug fix

* Model training example

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

apax/utils/__init__.pyi

@@ -1,3 +1,3 @@
-from . import convert, data, jax_md_reduced, math, random, datasets
+from . import convert, data, datasets, jax_md_reduced, math, random
 
 __all__ = ["convert", "data", "math", "random", "jax_md_reduced", datasets]

apax/utils/helpers.py

@@ -17,5 +17,8 @@ def mod_config(config_path, updated_config):
         config_dict = yaml.safe_load(stream)
 
     for key, new_value in updated_config.items():
-        config_dict[key].update(new_value)
+        if isinstance(config_dict[key], dict):
+            config_dict[key].update(new_value)
+        else:
+            config_dict[key] = new_value
     return config_dict

tests/conftest.py

@@ -12,9 +12,9 @@
 from apax.config.train_config import Config
 from apax.model.builder import ModelBuilder
 from apax.train.run import run
-from apax.utils.random import seed_py_np_tf
 from apax.utils.datasets import download_md22_stachyose
 from apax.utils.helpers import mod_config
+from apax.utils.random import seed_py_np_tf
 
 
 @pytest.fixture(autouse=True)