From 40e0b201e61fa0604fc8946450305dcba7f83e5d Mon Sep 17 00:00:00 2001 From: Daniel Rosel Date: Thu, 22 Jan 2026 13:10:15 +0100 Subject: [PATCH] chore: init code for jax core --- sim/rl/jax_core/__init__.py | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 sim/rl/jax_core/__init__.py diff --git a/sim/rl/jax_core/__init__.py b/sim/rl/jax_core/__init__.py new file mode 100644 index 0000000..99d5a87 --- /dev/null +++ b/sim/rl/jax_core/__init__.py @@ -0,0 +1,11 @@ +"""JAX-accelerated simulation core for PHANTOM environment.""" +from .transitions import TransitionData, compile_transitions, fallback_transitions, JAX_AVAILABLE +from .simulation import SessionBatch, SimResult, sample_sessions, compute_metrics +from .features import session_features, compute_session_transitions +from .separability import compute_divergences, estimate_alpha_batch + +__all__ = [ + "JAX_AVAILABLE", "TransitionData", "compile_transitions", "fallback_transitions", + "SessionBatch", "SimResult", "sample_sessions", "compute_metrics", + "session_features", "compute_session_transitions", "compute_divergences", "estimate_alpha_batch", +]