Towards Unified Sequence Representation IV: Add Unified_RNN for SequenceProcessorV2 fields

docs/api/models.rst

@@ -39,6 +39,9 @@ routes each feature type automatically.
    * - ``MultimodalRNN``
      - Features mix sequential codes with static tensors or multi-hot encodings
      - Auto-routes sequential features to RNN layers and non-sequential features to linear layers; good default for EHR
+   * - ``Unified_RNN``
+     - Features are ``sequence_v2`` fields (nested codes, values and timestamps), plus static tensors
+     - Encodes each field bottom-up over packed records, without padding the hierarchy
    * - :doc:`models/pyhealth.models.StageNet`
      - You have time-stamped vital signs with irregular measurement intervals
      - Requires ``StageNetProcessor`` or ``StageNetTensorProcessor`` in the task schema

docs/api/models/pyhealth.models.RNN.rst

@@ -17,4 +17,9 @@ The separate callable RNNLayer and the complete RNN model.
 .. autoclass:: pyhealth.models.MultimodalRNN
     :members:
     :undoc-members:
-    :show-inheritance:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.Unified_RNN
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst

@@ -5,6 +5,7 @@ Embeddings for the packed records of ``sequence_v2``
 (:class:`~pyhealth.processors.SequenceProcessorV2`) fields. Each level's
 records get one embedding per token namespace and a projection of their
 values, concatenated; each level keeps its offsets and times.
+:class:`~pyhealth.models.Unified_RNN` builds on it.
 
 .. autoclass:: pyhealth.models.SequenceEmbeddingModel
     :members:

pyhealth/models/__init__.py

@@ -28,6 +28,7 @@
 from .molerec import MoleRec, MoleRecLayer
 from .retain import MultimodalRETAIN, RETAIN, RETAINLayer
 from .rnn import MultimodalRNN, RNN, RNNLayer
+from .rnn import Unified_RNN as Unified_RNN
 from .safedrug import SafeDrug, SafeDrugLayer
 from .sparcnet import DenseBlock, DenseLayer, SparcNet, TransitionLayer
 from .stagenet import StageNet, StageNetLayer

pyhealth/models/base_model.py

@@ -46,7 +46,7 @@ def _label_record_loss(
         raise ValueError(
             f"Expected one prediction per label record, {tuple(values.shape)}, "
             f"got {tuple(logits.shape)}. Per-step labels need a model that "
-            "predicts at each label's time."
+            "predicts at each label's time, such as Unified_RNN."
         )
     return loss_fn(logits, values)