Refactor configs (#54)

---------

Co-authored-by: Charles Zhang <charleszhang@holylogin04.rc.fas.harvard.edu>

stac_mjx/io.py

@@ -12,9 +12,19 @@
 import h5py
 from pathlib import Path
 from omegaconf import DictConfig
+from jax.lib import xla_bridge
 import os
 
 
+def enable_xla_flags():
+    """Enables XLA Flags for faster runtime on Nvidia GPUs."""
+    if xla_bridge.get_backend().platform == "gpu":
+        os.environ["XLA_FLAGS"] = (
+            "--xla_gpu_enable_triton_softmax_fusion=true "
+            "--xla_gpu_triton_gemm_any=True "
+        )
+
+
 def load_data(cfg: DictConfig, base_path: Union[Path, None] = None):
     """Main mocap data file loader interface.