xax 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.1"
+__version__ = "0.3.3"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -42,11 +42,12 @@ __all__ = [
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "normalize",
+    "quat_mul",
     "quat_to_euler",
     "quat_to_rotmat",
-    "quat_mul",
     "rotate_vector_by_quat",
     "rotation6d_to_rotation_matrix",
+    "rotation_matrix_to_quat",
     "rotation_matrix_to_rotation6d",
     "cross_entropy",
     "cast_norm_type",
@@ -224,11 +225,12 @@ NAME_MAP: dict[str, str] = {
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "normalize": "nn.geom",
+    "quat_mul": "nn.geom",
     "quat_to_euler": "nn.geom",
     "quat_to_rotmat": "nn.geom",
-    "quat_mul": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
     "rotation6d_to_rotation_matrix": "nn.geom",
+    "rotation_matrix_to_quat": "nn.geom",
     "rotation_matrix_to_rotation6d": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.metrics",
@@ -405,6 +407,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         quat_to_rotmat,
         rotate_vector_by_quat,
         rotation6d_to_rotation_matrix,
+        rotation_matrix_to_quat,
         rotation_matrix_to_rotation6d,
     )
     from xax.nn.losses import cross_entropy

xax/core/state.py

@@ -107,7 +107,7 @@ class State:
     @classmethod
     def from_dict(cls, **d: Unpack[StateDict]) -> "State":
         if "phase" in d:
-            d["_phase"] = _phase_to_int(cast(Phase, d.pop("phase")))
+            d["_phase"] = _phase_to_int(d.pop("phase"))
 
         int32_arr = jnp.array(
             [

xax/nn/geom.py

@@ -272,3 +272,85 @@ def quat_mul(q2: Array, q1: Array) -> Array:
     z = w2 * z1 + x2 * y1 - y2 * x1 + z2 * w1
 
     return jnp.concatenate([w, x, y, z], axis=-1)
+
+
+def rotation_matrix_to_quat(rotation_matrix: Array, eps: float = 1e-6) -> Array:
+    """Converts a rotation matrix to a unit quaternion ``(w, x, y, z)``.
+
+    Args:
+        rotation_matrix: The rotation matrix, shape ``(*, 3, 3)``.
+        eps: A small epsilon value to avoid division by zero when normalising.
+
+    Returns:
+        A quaternion with shape ``(*, 4)``.
+    """
+    chex.assert_shape(rotation_matrix, (..., 3, 3))
+
+    m00 = rotation_matrix[..., 0, 0]
+    m01 = rotation_matrix[..., 0, 1]
+    m02 = rotation_matrix[..., 0, 2]
+    m10 = rotation_matrix[..., 1, 0]
+    m11 = rotation_matrix[..., 1, 1]
+    m12 = rotation_matrix[..., 1, 2]
+    m20 = rotation_matrix[..., 2, 0]
+    m21 = rotation_matrix[..., 2, 1]
+    m22 = rotation_matrix[..., 2, 2]
+
+    trace = m00 + m11 + m22
+
+    # Case 0: trace is positive
+    s0 = jnp.sqrt(jnp.clip(trace + 1.0, a_min=0.0)) * 2.0  # S = 4 * qw
+    w0 = 0.25 * s0
+    x0 = (m21 - m12) / jnp.where(s0 < eps, 1.0, s0)
+    y0 = (m02 - m20) / jnp.where(s0 < eps, 1.0, s0)
+    z0 = (m10 - m01) / jnp.where(s0 < eps, 1.0, s0)
+
+    # Case 1: m00 is the largest diagonal term
+    s1 = jnp.sqrt(jnp.clip(1.0 + m00 - m11 - m22, a_min=0.0)) * 2.0  # S = 4 * qx
+    w1 = (m21 - m12) / jnp.where(s1 < eps, 1.0, s1)
+    x1 = 0.25 * s1
+    y1 = (m01 + m10) / jnp.where(s1 < eps, 1.0, s1)
+    z1 = (m02 + m20) / jnp.where(s1 < eps, 1.0, s1)
+
+    # Case 2: m11 is the largest diagonal term
+    s2 = jnp.sqrt(jnp.clip(1.0 + m11 - m00 - m22, a_min=0.0)) * 2.0  # S = 4 * qy
+    w2 = (m02 - m20) / jnp.where(s2 < eps, 1.0, s2)
+    x2 = (m01 + m10) / jnp.where(s2 < eps, 1.0, s2)
+    y2 = 0.25 * s2
+    z2 = (m12 + m21) / jnp.where(s2 < eps, 1.0, s2)
+
+    # Case 3: m22 is the largest diagonal term
+    s3 = jnp.sqrt(jnp.clip(1.0 + m22 - m00 - m11, a_min=0.0)) * 2.0  # S = 4 * qz
+    w3 = (m10 - m01) / jnp.where(s3 < eps, 1.0, s3)
+    x3 = (m02 + m20) / jnp.where(s3 < eps, 1.0, s3)
+    y3 = (m12 + m21) / jnp.where(s3 < eps, 1.0, s3)
+    z3 = 0.25 * s3
+
+    cond0 = trace > 0.0
+    cond1 = (m00 > m11) & (m00 > m22)
+    cond2 = m11 > m22
+
+    w = jnp.where(
+        cond0,
+        w0,
+        jnp.where(cond1, w1, jnp.where(cond2, w2, w3)),
+    )
+    x = jnp.where(
+        cond0,
+        x0,
+        jnp.where(cond1, x1, jnp.where(cond2, x2, x3)),
+    )
+    y = jnp.where(
+        cond0,
+        y0,
+        jnp.where(cond1, y1, jnp.where(cond2, y2, y3)),
+    )
+    z = jnp.where(
+        cond0,
+        z0,
+        jnp.where(cond1, z1, jnp.where(cond2, z2, z3)),
+    )
+
+    quat = jnp.stack([w, x, y, z], axis=-1)
+    quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
+    return quat

xax/task/base.py

@@ -82,6 +82,9 @@ class BaseTask(Generic[Config]):
     def on_after_checkpoint_save(self, ckpt_path: Path, state: State | None) -> State | None:
         return state
 
+    def add_logger_handlers(self, logger: logging.Logger) -> None:
+        pass
+
     @functools.cached_property
     def task_class_name(self) -> str:
         return self.__class__.__name__

xax/task/launchers/single_process.py

@@ -15,8 +15,9 @@ def run_single_process_training(
     *cfgs: RawConfigType,
     use_cli: bool | list[str] = True,
 ) -> None:
-    configure_logging()
+    logger = configure_logging()
     task_obj = task.get_task(*cfgs, use_cli=use_cli)
+    task_obj.add_logger_handlers(logger)
     task_obj.run()
 
 

xax/task/mixins/artifacts.py

@@ -14,7 +14,7 @@ from xax.core.state import State
 from xax.nn.parallel import is_master
 from xax.task.base import BaseConfig, BaseTask
 from xax.utils.experiments import stage_environment
-from xax.utils.logging import LOG_STATUS
+from xax.utils.logging import LOG_STATUS, RankFilter
 from xax.utils.text import show_info
 
 logger = logging.getLogger(__name__)
@@ -24,6 +24,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class ArtifactsConfig(BaseConfig):
     exp_dir: str | None = field(None, help="The fixed experiment directory")
+    log_to_file: bool = field(True, help="If set, add a file handler to the logger to write all logs to the exp dir")
 
 
 Config = TypeVar("Config", bound=ArtifactsConfig)
@@ -39,6 +40,14 @@ class ArtifactsMixin(BaseTask[Config]):
         self._exp_dir = None
         self._stage_dir = None
 
+    def add_logger_handlers(self, logger: logging.Logger) -> None:
+        super().add_logger_handlers(logger)
+        if is_master() and self.config.log_to_file:
+            file_handler = logging.FileHandler(self.exp_dir / "logs.txt")
+            file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+            file_handler.addFilter(RankFilter(rank=0))
+            logger.addHandler(file_handler)
+
     @functools.cached_property
     def run_dir(self) -> Path:
         run_dir = get_run_dir()

xax/task/mixins/runnable.py

@@ -10,6 +10,7 @@ import jax
 
 from xax.task.base import BaseConfig, BaseTask, RawConfigType
 from xax.task.launchers.base import BaseLauncher
+from xax.task.launchers.cli import CliLauncher
 
 
 @jax.tree_util.register_dataclass
@@ -45,8 +46,6 @@ class RunnableMixin(BaseTask[Config], ABC):
         use_cli: bool | list[str] = True,
     ) -> None:
         if launcher is None:
-            from xax.task.launchers.cli import CliLauncher
-
             launcher = CliLauncher()
         launcher.launch(cls, *cfgs, use_cli=use_cli)
 

xax/utils/logging.py

@@ -146,7 +146,7 @@ def configure_logging(
     rank: int | None = None,
     world_size: int | None = None,
     debug: bool | None = None,
-) -> None:
+) -> logging.Logger:
     """Instantiates logging.
 
     This captures logs and reroutes them to the Toasts module, which is
@@ -186,6 +186,8 @@ def configure_logging(
         logging.getLogger("PIL").setLevel(logging.WARNING)
         logging.getLogger("torch").setLevel(logging.WARNING)
 
+    return root_logger
+
 
 def get_unused_port(default: int | None = None) -> int:
     """Returns an unused port number on the local machine.

{xax-0.3.1.dist-info → xax-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.3.1
+Version: 0.3.3
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.3.1.dist-info → xax-0.3.3.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=Ku51ccrDjxWuSCgmHHVnTXhloyA_d9A1VlEEbO0Ycjg,15713
+xax/__init__.py,sha256=ffVd9_qSVuEAIPn6eK_6N8qEfDALRJigArHZQGy6y1o,15819
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
@@ -6,25 +6,25 @@ xax/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/cli/edit_config.py,sha256=LQUIlOS6hvPZyVEaMme3FP-62M0BKQPYavCwVDWuBLw,2600
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
-xax/core/state.py,sha256=F9Tj3FfCw8zFKaDEoEGiThZE2ntYEtzNjnBX3pQ1g60,3826
+xax/core/state.py,sha256=_gtINsRc310Bu_HuIYsDoOKTZa6DgU2tz0IOKkdnY9Q,3813
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/attention.py,sha256=0essK90OO3x9FxnUqU0DhufwXKRMN41zMtRCki5iKzQ,24742
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
-xax/nn/geom.py,sha256=7_cUeoJIfCS_YbQ6O7zElnpYsAn4XrVm4m7PaT9UcYw,8005
+xax/nn/geom.py,sha256=6rBQrZRX1miG08VG-s8phPjA6MEFxUAfQVPt5F0RQQI,10645
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
 xax/nn/metrics.py,sha256=zuvPXlRQczBTLHD4ilNGmZaiq6Yie3rxCMq6JkI_kos,3154
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=TYANmjNcce4_V5ZSYLnE91PXRn7Nn0nT7hN8plW_Au0,8117
+xax/task/base.py,sha256=i6FRJ75aqlekWkzJNRWDUEX7P514pUjLVuxjhX1GBgw,8198
 xax/task/logger.py,sha256=Bmhl4mv08Aq49ZyX6BdjPIsPJK28e8s3mVFatM4IY2Q,41060
 xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
 xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
 xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,1402
-xax/task/launchers/single_process.py,sha256=IoML-30g5c526yxkpbWSOtG_KpNQMakT7xujzB1gIAo,846
+xax/task/launchers/single_process.py,sha256=wdEUT-B-FE9aemmt1tB_rKKRNy60aiDhslsy2i-ojWo,896
 xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/loggers/callback.py,sha256=zQuV1xCvz47Q3UQqP1D5mBhbVzptvmPR_7hX25vqSk0,1667
 xax/task/loggers/json.py,sha256=6A5wL7kspsXnpPhI_vu0scgd2Z2-WLhw4gbBFm7eZMM,4377
@@ -32,7 +32,7 @@ xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,140
 xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
 xax/task/loggers/tensorboard.py,sha256=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
-xax/task/mixins/artifacts.py,sha256=IBPAQMCGd7PQiZHfSjLakPW5j7cNuL6AsW6QkVSc02E,3277
+xax/task/mixins/artifacts.py,sha256=R-y3p7__zJHlHDqwDVAZysg2ZmebCJbqAx_xGT2Xpd0,3857
 xax/task/mixins/checkpointing.py,sha256=v50IZ7j58DWmEu-_6Zh_02R5KUVGhrMkg5n-MYM_J4c,11484
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
@@ -40,7 +40,7 @@ xax/task/mixins/data_loader.py,sha256=Tp7zqPdfH2_JuE6J6EP-fEtCQpq9MjKlGHYK7Zh-go
 xax/task/mixins/gpu_stats.py,sha256=USOyhXldxbsrl6eCtoFKTWUm_lfeG0cUCkQNUpXRdtA,8880
 xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,2808
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
-xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
+xax/task/mixins/runnable.py,sha256=pcLrYc_TycZUY9zZim05Skc2FWk3IZKFnu6p3UDMonM,1966
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
 xax/task/mixins/train.py,sha256=TZatz5QwTfrNhQTiO2IqrmQY9P4Lay6FAD2VsQpWa54,33245
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -48,7 +48,7 @@ xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=5k5hPYSaVjzoR_nm2Q3DAHMMYi3Bcp3N3PAQbwZq7Gg,29830
 xax/utils/jax.py,sha256=6cP95-rcjkRt1fefkZWJQhJhH0uUYWJB3w4NP1-aDp0,10136
 xax/utils/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
-xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
+xax/utils/logging.py,sha256=Kkyma_LJXqrN2HTQ214gRP_9ih3_bKk115MWC60lQWM,6656
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
 xax/utils/pytree.py,sha256=rVY2kKa637xfX3Oue6OP9ScwmDyxJ_CeHkUpZZtmN04,9231
@@ -59,9 +59,9 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.3.1.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.3.1.dist-info/METADATA,sha256=xAwDUWNoga60_KZ_nPEXrdAuvlak-zWAXk3lE-XvLeI,1246
-xax-0.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.3.1.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.3.1.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.3.1.dist-info/RECORD,,
+xax-0.3.3.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.3.dist-info/METADATA,sha256=mjIzoFZDSR3V1-2LHbvup6wDVa4vLiqbqiNWLsKCXY8,1246
+xax-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.3.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.3.dist-info/RECORD,,