Use marathon-train from PyPI and clean up lazy imports

Switch marathon-train dependency from pinned git URL to PyPI package and move marathon imports to module level in backbone.py and calculator.py (the lazy pattern was a workaround for the git dependency).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: sirmarcel

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
     "jax",
     "jax-pme @ git+https://github.com/lab-cosmo/jax-pme.git",
     "jaxtyping",
-    "marathon-train[grain] @ git+https://github.com/sirmarcel/marathon.git@release/v0.2.0",
+    "marathon-train[grain]",
     "opsis",
     "optax",
 ]

src/lorem/calculator.py

@@ -8,6 +8,8 @@
     BaseCalculator,
     PropertyNotImplementedError,
 )
+from marathon.emit.checkpoint import read_msgpack
+from marathon.io import from_dict, read_yaml
 
 from lorem.neighborlist import NeighborListCache
 
@@ -78,8 +80,6 @@ def from_checkpoint(
     ):
         from pathlib import Path
 
-        from marathon.io import from_dict, read_yaml
-
         folder = Path(folder)
 
         model = from_dict(read_yaml(folder / "model/model.yaml"))
@@ -89,8 +89,6 @@ def from_checkpoint(
         baseline = read_yaml(folder / "model/baseline.yaml")
         species_to_weight = baseline["elemental"]
 
-        from marathon.emit.checkpoint import read_msgpack
-
         params = read_msgpack(folder / "model/model.msgpack")
 
         return cls(model.predict, species_to_weight, params, model.cutoff, **kwargs)

src/lorem/models/backbone.py

@@ -8,12 +8,11 @@
 import e3x
 import flax.linen as nn
 from flax.core import FrozenDict
+from marathon.utils import masked
 
 
 def _masked(fn, x, mask):
-    """Apply fn only where mask is True. Lazy import from marathon."""
-    from marathon.utils import masked
-
+    """Apply fn only where mask is True."""
     return masked(fn, x, mask)