xax 0.3.0__tar.gz → 0.3.2__tar.gz

{xax-0.3.0/xax.egg-info → xax-0.3.2}/PKG-INFO

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

{xax-0.3.0 → xax-0.3.2}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.2"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -44,6 +44,7 @@ __all__ = [
     "normalize",
     "quat_to_euler",
     "quat_to_rotmat",
+    "quat_mul",
     "rotate_vector_by_quat",
     "rotation6d_to_rotation_matrix",
     "rotation_matrix_to_rotation6d",
@@ -225,6 +226,7 @@ NAME_MAP: dict[str, str] = {
     "normalize": "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_rotation6d": "nn.geom",
@@ -398,6 +400,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
         normalize,
+        quat_mul,
         quat_to_euler,
         quat_to_rotmat,
         rotate_vector_by_quat,

{xax-0.3.0 → xax-0.3.2}/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-0.3.0 → xax-0.3.2}/xax/nn/geom.py

@@ -251,3 +251,24 @@ def rotation_matrix_to_rotation6d(rotation_matrix: jnp.ndarray) -> jnp.ndarray:
     # Simply concatenate a1 and a2 from SO(3)
     r6d = jnp.concatenate([rotation_matrix[..., 0], rotation_matrix[..., 1]], axis=-1)
     return r6d.reshape(shape[:-2] + (6,))
+
+
+def quat_mul(q2: Array, q1: Array) -> Array:
+    """Multiply two quaternions (supports batching).
+
+    Args:
+        q2: Second quaternion (w, x, y, z), shape (..., 4)
+        q1: First quaternion (w, x, y, z), shape (..., 4)
+
+    Returns:
+        Product quaternion, shape (..., 4)
+    """
+    w1, x1, y1, z1 = jnp.split(q1, 4, axis=-1)
+    w2, x2, y2, z2 = jnp.split(q2, 4, axis=-1)
+
+    w = w2 * w1 - x2 * x1 - y2 * y1 - z2 * z1
+    x = w2 * x1 + x2 * w1 + y2 * z1 - z2 * y1
+    y = w2 * y1 - x2 * z1 + y2 * w1 + z2 * x1
+    z = w2 * z1 + x2 * y1 - y2 * x1 + z2 * w1
+
+    return jnp.concatenate([w, x, y, z], axis=-1)

{xax-0.3.0 → xax-0.3.2}/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-0.3.0 → xax-0.3.2}/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-0.3.0 → xax-0.3.2}/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-0.3.0 → xax-0.3.2}/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.0 → xax-0.3.2/xax.egg-info}/PKG-INFO

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