[BMS] StageNet with Attention

docs/api/models.rst

@@ -24,6 +24,7 @@ We implement the following models for supporting multiple healthcare predictive
     models/pyhealth.models.ContraWR
     models/pyhealth.models.SparcNet
     models/pyhealth.models.StageNet
+    models/pyhealth.models.StageAttentionNet
     models/pyhealth.models.AdaCare
     models/pyhealth.models.ConCare
     models/pyhealth.models.Agent

docs/api/models/pyhealth.models.StageAttentionNet.rst

@@ -0,0 +1,14 @@
+pyhealth.models.StageAttentionNet
+===================================
+
+The separate callable StageNetAttentionLayer and the complete StageAttentionNet model.
+
+.. autoclass:: pyhealth.models.StageNetAttentionLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.StageAttentionNet
+    :members:
+    :undoc-members:
+    :show-inheritance:

pyhealth/models/__init__.py

@@ -21,6 +21,7 @@
 from .safedrug import SafeDrug, SafeDrugLayer
 from .sparcnet import DenseBlock, DenseLayer, SparcNet, TransitionLayer
 from .stagenet import StageNet, StageNetLayer
+from .stagenet_mha import StageAttentionNet, StageNetAttentionLayer
 from .tcn import TCN, TCNLayer
 from .torchvision_model import TorchvisionModel
 from .transformer import Transformer, TransformerLayer