refactoring training spc setup and benchmarking

engine/backends/jax.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+from ..jax import JAX_AVAILABLE
+
+
+def train_jax_backend(
+    cfg: Mapping[str, Any],
+) -> tuple[dict[str, Any], dict[str, float | int | str]]:
+    if not JAX_AVAILABLE:
+        raise ImportError(
+            "JAX backend requested but JAX is not installed. "
+            "Install engine/jax/requirements.txt and jax[tpu] for TPU runs."
+        )
+    from ..jax.train import train_jax
+
+    return train_jax(dict(cfg))