chore: jax<0.7.0, other minor code cleanup

Author: ttngu207

CHANGELOG.md

@@ -3,6 +3,10 @@
 Observes [Semantic Versioning](https://semver.org/spec/v2.0.0.html) standard and
 [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) convention.
 
+## [1.0.2] - 2025-10-07
++ Update - `keypoint-moseq` as extra dependency
++ Fix - Version pin `jax<0.7.0`
+
 ## [1.0.1] - 2025-09-23
 + Feat - Add support to generate PNG version of fitting progress plots in `PreFit`, `FullFit`, and `moseq_report` schema
 + Fix - Update path handling to use `Path` objects and `dj.logger`

README.md

@@ -25,7 +25,7 @@ DataJoint Elements collectively standardize and automate data collection and ana
 + Install with `pip`:
 
   ```bash
-  pip install -e .
+  pip install -e .[keypoint-moseq]
   ```
 
 + [Interactive tutorial on GitHub Codespaces](https://github.com/datajoint/element-moseq#interactive-tutorial)

element_moseq/plotting/viz_utils.py

@@ -10,9 +10,6 @@
 import datajoint as dj
 import matplotlib.pyplot as plt
 import numpy as np
-from jax_moseq.models.keypoint_slds import center_embedding
-from keypoint_moseq.util import get_distance_to_medoid, get_edges, plot_keypoint_traces
-from keypoint_moseq.viz import plot_pcs_3D
 
 logger = dj.logger
 
@@ -132,6 +129,7 @@ def plot_medoid_distance_outliers(
     None
         The plot is saved to 'QA/plots/keypoint_distance_outliers/{recording_name}.png'.
     """
+    from keypoint_moseq.util import get_distance_to_medoid, plot_keypoint_traces
 
     plot_path = os.path.join(
         project_dir,
@@ -241,6 +239,10 @@ def plot_pcs(
     interactive : bool, default=True
         For 3D data, whether to generate an interactive 3D plot.
     """
+    from jax_moseq.models.keypoint_slds import center_embedding
+    from keypoint_moseq.util import get_edges
+    from keypoint_moseq.viz import plot_pcs_3D
+
     k = len(use_bodyparts)
     d = len(pca.mean_) // (k - 1)
 

element_moseq/readers/kpms_reader.py

@@ -3,7 +3,6 @@
 from typing import Any, Dict, Union
 
 import datajoint as dj
-import jax.numpy as jnp
 import yaml
 
 logger = dj.logger
@@ -103,6 +102,8 @@ def load_kpms_dj_config(
       - build_indexes -> adds jax arrays 'anterior_idxs' and 'posterior_idxs'
         indexing into 'use_bodyparts' by order.
     """
+    import jax.numpy as jnp
+
     dj_cfg_path = _dj_config_path(project_dir)
     if not Path(dj_cfg_path).exists():
         raise FileNotFoundError(

element_moseq/version.py

@@ -2,4 +2,4 @@
 Package metadata
 """
 
-__version__ = "1.0.1"
+__version__ = "1.0.2"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "element-moseq"
-version = "1.0.1"
+version = "1.0.2"
 description = "Keypoint-MoSeq DataJoint Element"
 readme = "README.md"
 license = {text = "MIT"}
@@ -24,7 +24,6 @@ dependencies = [
     "ipykernel>=6.0.1",
     "ipywidgets",
     "opencv-python",
-    "keypoint-moseq @ git+https://github.com/dattalab/keypoint-moseq/",
     "pdf2image",
 ]
 
@@ -35,6 +34,7 @@ elements = [
     "element-interface @ git+https://github.com/datajoint/element-interface.git",
     "element-animal @ git+https://github.com/datajoint/element-animal.git",
 ]
+keypoint-moseq = ["jax<0.7.0", "keypoint-moseq @ git+https://github.com/dattalab/keypoint-moseq.git"]
 tests = [
     "pytest",
     "pytest-cov",