xax 0.2.8__tar.gz → 0.2.10__tar.gz

{xax-0.2.8/xax.egg-info → xax-0.2.10}/PKG-INFO

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

{xax-0.2.8 → xax-0.2.10}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.8"
+__version__ = "0.2.10"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -308,7 +308,7 @@ NAME_MAP: dict[str, str] = {
     "compute_nan_ratio": "utils.pytree",
     "flatten_array": "utils.pytree",
     "flatten_pytree": "utils.pytree",
-    "get_param_count": "utils.pytree",
+    "get_pytree_param_count": "utils.pytree",
     "pytree_has_nans": "utils.pytree",
     "reshuffle_pytree": "utils.pytree",
     "reshuffle_pytree_along_dims": "utils.pytree",

{xax-0.2.8 → xax-0.2.10}/xax/task/mixins/train.py

@@ -361,7 +361,6 @@ class TrainMixin(
         model = self.get_model(key)
         state = State.init_state()
 
-        self.log_model_size(model)
         if not load_optimizer:
             return model, state
 
@@ -684,9 +683,6 @@ class TrainMixin(
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
         self.logger.log_file("info.json", get_info_json())
 
-    def log_model_size(self, model: PyTree) -> None:
-        logger.info("Model size: %s", f"{get_pytree_param_count(model):,}")
-
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
         return eqx.is_inexact_array(item)
 
@@ -832,6 +828,9 @@ class TrainMixin(
 
             key, model_key = jax.random.split(key)
             model, optimizer, opt_state, state = self.load_initial_state(model_key, load_optimizer=True)
+            logger.info("Model size: %s", f"{get_pytree_param_count(model):,}")
+            logger.info("Optimizer size: %s", f"{get_pytree_param_count(optimizer):,}")
+
             state = self.on_training_start(state)
 
             def on_exit() -> None:

{xax-0.2.8 → xax-0.2.10/xax.egg-info}/PKG-INFO

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