DeepHealth/models.py

from collections.abc import Mapping
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from backbones import (
    AgeSinusoidalEncoding,
    GaussianRBFTimeBasis,
    SharedEventTrajectoryCore,
    TimeRoPE,
    TokenAutoDiscretization,
)
from targets import PAD_IDX


EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2"


@dataclass(frozen=True)
class EventTrajectoryModelSize:
    d_model: int
    n_trajectory: int

    @property
    def trajectory_dim(self) -> int:
        return self.d_model // self.n_trajectory

    @property
    def traj_hidden(self) -> int:
        return 4 * self.n_trajectory


MODEL_SIZE_PRESETS = {
    "nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6),
    "tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
    "small": EventTrajectoryModelSize(d_model=512, n_trajectory=8),
    "medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12),
    "huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16),
}
MODEL_SIZE_NAMES = tuple(MODEL_SIZE_PRESETS)


def resolve_model_size(model_size: str) -> EventTrajectoryModelSize:
    if not isinstance(model_size, str):
        raise ValueError(
            f"model_size must be a string, got {type(model_size).__name__}"
        )
    normalized = model_size.strip().lower()
    try:
        return MODEL_SIZE_PRESETS[normalized]
    except KeyError as exc:
        choices = ", ".join(MODEL_SIZE_NAMES)
        raise ValueError(
            f"Unknown model_size {model_size!r}; expected one of: {choices}"
        ) from exc


def _required_config_int(
    config: Mapping[str, object],
    key: str,
) -> int:
    raw_value = config.get(key)
    if isinstance(raw_value, bool):
        raise ValueError(f"Config field {key!r} must be an integer")
    try:
        value = int(raw_value)
    except (TypeError, ValueError) as exc:
        raise ValueError(
            f"Config field {key!r} must be present and integer-valued; "
            f"got {raw_value!r}"
        ) from exc
    if isinstance(raw_value, float) and not raw_value.is_integer():
        raise ValueError(f"Config field {key!r} must be an integer")
    return value


def validate_event_trajectory_config(config: Mapping[str, object]) -> None:
    actual = config.get("model_architecture")
    if actual != EVENT_TRAJECTORY_ARCHITECTURE:
        raise ValueError(
            "This branch only accepts models trained with the shared "
            "event-trajectory architecture marker "
            f"{EVENT_TRAJECTORY_ARCHITECTURE!r}; got {actual!r}."
        )
    raw_model_size = config.get("model_size")
    if not isinstance(raw_model_size, str):
        raise ValueError(
            "Config field 'model_size' must be one of: "
            + ", ".join(MODEL_SIZE_NAMES)
        )
    model_size = raw_model_size.strip().lower()
    preset = resolve_model_size(model_size)
    d_model = _required_config_int(config, "d_model")
    n_trajectory = _required_config_int(config, "n_trajectory")
    n_reasoning_rounds = _required_config_int(
        config,
        "n_reasoning_rounds",
    )
    trajectory_dim = _required_config_int(config, "trajectory_dim")
    traj_hidden = _required_config_int(config, "traj_hidden")
    if n_reasoning_rounds <= 0:
        raise ValueError(
            "n_reasoning_rounds must be positive"
        )
    expected_values = {
        "d_model": preset.d_model,
        "n_trajectory": preset.n_trajectory,
        "trajectory_dim": preset.trajectory_dim,
        "traj_hidden": preset.traj_hidden,
    }
    actual_values = {
        "d_model": d_model,
        "n_trajectory": n_trajectory,
        "trajectory_dim": trajectory_dim,
        "traj_hidden": traj_hidden,
    }
    mismatches = [
        f"{key}: expected {expected}, got {actual_values[key]}"
        for key, expected in expected_values.items()
        if actual_values[key] != expected
    ]
    if mismatches:
        raise ValueError(
            f"Config does not match model_size={model_size!r}: "
            + "; ".join(mismatches)
        )


def _checkpoint_scalar_int(
    state_dict: Mapping[str, object],
    key: str,
) -> int:
    value = state_dict[key]
    if not isinstance(value, torch.Tensor) or value.numel() != 1:
        raise ValueError(
            f"Checkpoint architecture field {key!r} must be a scalar tensor"
        )
    return int(value.detach().cpu().item())


def validate_event_trajectory_state_dict(
    state_dict: Mapping[str, object],
    *,
    expected_d_model: int | None = None,
    expected_n_trajectory: int | None = None,
    expected_n_reasoning_rounds: int | None = None,
) -> None:
    required_keys = {
        "architecture_d_model",
        "architecture_n_trajectory",
        "architecture_n_reasoning_rounds",
        "event_projection.weight",
        "trajectory_prototypes",
        "query_projection.weight",
        "reasoning_core.cross_attention.q_proj.weight",
        "reasoning_core.cross_attention.k_proj.weight",
        "reasoning_core.cross_attention.v_proj.weight",
        "reasoning_core.traj_mixer.gate_proj",
        "reasoning_core.traj_mixer.value_proj",
        "reasoning_core.traj_mixer.output_proj",
        "reasoning_core.attn_scale",
        "reasoning_core.mixer_scale",
    }
    missing = sorted(required_keys.difference(state_dict))
    if missing:
        raise ValueError(
            "Checkpoint is not a shared event-trajectory checkpoint; "
            "missing required "
            f"parameters: {', '.join(missing)}"
        )
    checkpoint_values = {
        "d_model": _checkpoint_scalar_int(
            state_dict,
            "architecture_d_model",
        ),
        "n_trajectory": _checkpoint_scalar_int(
            state_dict,
            "architecture_n_trajectory",
        ),
        "n_reasoning_rounds": _checkpoint_scalar_int(
            state_dict,
            "architecture_n_reasoning_rounds",
        ),
    }
    expected_values = {
        "d_model": expected_d_model,
        "n_trajectory": expected_n_trajectory,
        "n_reasoning_rounds": expected_n_reasoning_rounds,
    }
    mismatches = [
        f"{name}: checkpoint={checkpoint_values[name]}, expected={expected}"
        for name, expected in expected_values.items()
        if expected is not None and checkpoint_values[name] != expected
    ]
    if mismatches:
        raise ValueError(
            "Checkpoint architecture does not match the constructed model: "
            + "; ".join(mismatches)
        )


@dataclass
class DeepHealthOutput:
    hidden: torch.Tensor
    time_seq: torch.Tensor
    padding_mask: torch.Tensor
    event_len: int


class OtherInfoTokenizer(nn.Module):
    PAD_KIND = 0
    CONT_KIND = 1
    CATE_KIND = 2

    def __init__(
        self,
        n_embd: int,
        n_types: int,
        n_cont_types: int,
        n_categories: int,
        cont_type_ids: list[int],
        n_value_kinds: int = 3,
        n_bins: int = 16,
    ):
        super().__init__()
        if len(cont_type_ids) != n_cont_types:
            raise ValueError(
                "cont_type_ids length must match n_cont_types, got "
                f"{len(cont_type_ids)} vs {n_cont_types}"
            )
        if n_types <= 0:
            raise ValueError(f"n_types must include PAD and be > 0, got {n_types}")
        if n_categories <= 0:
            raise ValueError(
                f"n_categories must include PAD and be > 0, got {n_categories}"
            )
        if n_value_kinds <= self.CATE_KIND:
            raise ValueError(
                f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
            )

        self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
        self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
        self.cont_value_encoder = (
            TokenAutoDiscretization(
                n_cont_types=n_cont_types,
                n_bins=n_bins,
                n_embd=n_embd,
            )
            if n_cont_types > 0
            else None
        )
        self.cate_value_emb = nn.Embedding(
            n_categories,
            n_embd,
            padding_idx=0,
        )

        cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
        for idx, type_id in enumerate(cont_type_ids):
            if type_id <= 0 or type_id >= n_types:
                raise ValueError(
                    f"continuous type id {type_id} must be in [1, {n_types})"
                )
            cont_type_index[type_id] = idx
        self.register_buffer(
            "cont_type_index",
            cont_type_index,
            persistent=False,
        )
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.type_emb.weight[0])
        nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.kind_emb.weight[0])
        nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.cate_value_emb.weight[0])

    def forward(
        self,
        other_type: torch.LongTensor,
        other_value: torch.Tensor,
        other_value_kind: torch.LongTensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if other_type.shape != other_value.shape:
            raise ValueError(
                "other_type and other_value must have the same shape, got "
                f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}"
            )
        if other_type.shape != other_value_kind.shape:
            raise ValueError(
                "other_type and other_value_kind must have the same shape, got "
                f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}"
            )

        other_valid = other_type > 0
        type_emb = self.type_emb(other_type)
        kind_emb = self.kind_emb(other_value_kind)
        value_emb = torch.zeros_like(type_emb)

        cont_pos = other_valid & (other_value_kind == self.CONT_KIND)
        if cont_pos.any():
            if self.cont_value_encoder is None:
                raise ValueError("continuous tokens found but n_cont_types is 0")
            cont_idx = self.cont_type_index[other_type[cont_pos]]
            if (cont_idx < 0).any():
                bad_type = other_type[cont_pos][cont_idx < 0][0].item()
                raise ValueError(
                    f"type_id={bad_type} is marked continuous but is not in "
                    "cont_type_ids"
                )
            value_emb[cont_pos] = self.cont_value_encoder(
                cont_type_idx=cont_idx,
                value=other_value[cont_pos].to(type_emb.dtype),
            )

        cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
        if cate_pos.any():
            cate_id = other_value[cate_pos].long()
            value_emb[cate_pos] = self.cate_value_emb(cate_id)

        out = type_emb + kind_emb + value_emb
        out = out * other_valid.unsqueeze(-1).to(out.dtype)
        return out, other_valid


class DeepHealth(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        model_size: str,
        n_reasoning_rounds: int,
        n_types: int,
        n_cont_types: int,
        n_categories: int,
        cont_type_ids: list[int],
        n_value_kinds: int = 3,
        n_bins: int = 16,
        target_mode: str = "next_token",  # "next_token" or "all_future"
        time_mode: str = "relative",  # "relative" or "absolute"
        dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
        extra_pool_reduce: str = "mean",
        dropout: float = 0.0,
    ):
        super().__init__()
        if target_mode not in ["next_token", "all_future"]:
            raise ValueError(
                "target_mode must be either 'next_token' or 'all_future'")
        if time_mode not in ["relative", "absolute"]:
            raise ValueError(
                "time_mode must be either 'relative' or 'absolute'")
        if dist_mode not in ["exponential", "weibull", "mixed"]:
            raise ValueError(
                "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
        if extra_pool_reduce not in {"mean", "sum"}:
            raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
        if n_reasoning_rounds <= 0:
            raise ValueError(
                "n_reasoning_rounds must be positive, got "
                f"{n_reasoning_rounds}"
            )
        size_config = resolve_model_size(model_size)
        normalized_model_size = model_size.strip().lower()
        d_model = size_config.d_model
        n_trajectory = size_config.n_trajectory

        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.gender_embedding = nn.Embedding(
            2, d_model)  # Assuming binary gender
        self.tokenizer = OtherInfoTokenizer(
            n_embd=d_model,
            n_types=n_types,
            n_cont_types=n_cont_types,
            n_categories=n_categories,
            cont_type_ids=cont_type_ids,
            n_value_kinds=n_value_kinds,
            n_bins=n_bins,
        )
        self.target_mode = target_mode
        self.time_mode = time_mode
        self.dist_mode = dist_mode
        self.extra_pool_reduce = extra_pool_reduce
        self.model_size = normalized_model_size
        self.d_model = d_model
        self.n_trajectory = n_trajectory
        self.trajectory_dim = d_model // n_trajectory
        self.traj_hidden = 4 * n_trajectory
        self.n_reasoning_rounds = n_reasoning_rounds
        self.vocab_size = vocab_size
        self.register_buffer(
            "architecture_d_model",
            torch.tensor(d_model, dtype=torch.int64),
        )
        self.register_buffer(
            "architecture_n_trajectory",
            torch.tensor(n_trajectory, dtype=torch.int64),
        )
        self.register_buffer(
            "architecture_n_reasoning_rounds",
            torch.tensor(n_reasoning_rounds, dtype=torch.int64),
        )
        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.token_embedding.weight[0])
        nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
        if dist_mode == "weibull":
            self.rho_head = nn.Linear(d_model, vocab_size)
            nn.init.zeros_(self.rho_head.weight)
            nn.init.constant_(self.rho_head.bias, 0.5413)

        if dist_mode == "mixed":
            self.death_idx = vocab_size - 1
            self.rho_death_head = nn.Linear(d_model, 1)
            nn.init.zeros_(self.rho_death_head.weight)
            nn.init.constant_(self.rho_death_head.bias, 0.5413)

        # Event and query time are encoded once before shared reasoning. In
        # relative mode, cross-attention additionally uses TimeRoPE and RBF.
        self.age_encoding = AgeSinusoidalEncoding(d_model)
        self.event_projection = nn.Linear(d_model, d_model, bias=False)
        self.event_norm = nn.LayerNorm(d_model)
        self.query_projection = nn.Linear(d_model, d_model, bias=False)
        self.trajectory_prototypes = nn.Parameter(
            torch.empty(n_trajectory, self.trajectory_dim)
        )
        self.query_token = nn.Parameter(torch.empty(d_model))
        nn.init.normal_(self.event_projection.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.query_projection.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.trajectory_prototypes, mean=0.0, std=0.02)
        nn.init.normal_(self.query_token, mean=0.0, std=0.02)

        use_relative_time = time_mode == "relative"
        self.reasoning_core = SharedEventTrajectoryCore(
            d_model=d_model,
            n_trajectory=n_trajectory,
            n_reasoning_rounds=n_reasoning_rounds,
            dropout=dropout,
            n_rbf_bases=16,
            use_time_rope=use_relative_time,
            use_rbf_bias=use_relative_time,
        )
        if use_relative_time:
            self.rope = TimeRoPE(self.trajectory_dim)
            self.rbf = GaussianRBFTimeBasis(
                n_bases=16,
                max_time_diff=40.0,
            )
        else:
            self.rope = None
            self.rbf = None

        self.final_ln = nn.LayerNorm(d_model)
        self.risk_head = nn.Linear(d_model, vocab_size, bias=False)
        if target_mode == "next_token":
            self.risk_head.weight = self.token_embedding.weight

    def _make_event_invalid_mask(
        self,
        event_valid_mask: torch.Tensor,
        event_time: torch.Tensor,
        query_time: torch.Tensor,
        query_position: torch.Tensor | None = None,
    ) -> torch.Tensor:
        valid_key = event_valid_mask[:, None, :]
        key_time = event_time[:, None, :]
        query_time = query_time[:, :, None]
        if query_position is None:
            visible_by_time = key_time <= query_time
        else:
            key_position = torch.arange(
                event_time.size(1),
                device=event_time.device,
            ).view(1, 1, -1)
            visible_by_time = (key_time < query_time) | (
                (key_time == query_time)
                & (key_position <= query_position[:, :, None])
            )
        return ~(valid_key & visible_by_time)

    def _pool_other_by_time(
        self,
        h_other: torch.Tensor,
        other_time: torch.Tensor,
        other_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        batch_size, n_other, n_embd = h_other.shape
        if n_other == 0:
            empty_h = h_other.new_zeros(batch_size, 0, n_embd)
            empty_t = other_time.new_zeros(batch_size, 0)
            empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device)
            return empty_h, empty_t, empty_m

        masked_time = other_time.masked_fill(~other_mask, float("inf"))
        _sorted_time_with_pad, order = masked_time.sort(dim=1)
        sorted_time = other_time.gather(1, order)
        sorted_mask = other_mask.gather(1, order)
        sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd))

        group_start = torch.zeros_like(sorted_mask)
        group_start[:, 0] = sorted_mask[:, 0]
        group_start[:, 1:] = sorted_mask[:, 1:] & (
            sorted_time[:, 1:] != sorted_time[:, :-1]
        )
        group_id = group_start.long().cumsum(dim=1) - 1
        max_groups = int(group_start.sum(dim=1).max().item())

        pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
        pooled_time = other_time.new_zeros(batch_size, max_groups)
        pooled_mask = torch.zeros(
            batch_size,
            max_groups,
            dtype=torch.bool,
            device=h_other.device,
        )
        if max_groups == 0:
            return pooled_h, pooled_time, pooled_mask

        safe_group_id = group_id.clamp_min(0)
        pooled_h.scatter_add_(
            1,
            safe_group_id.unsqueeze(-1).expand_as(sorted_h),
            sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype),
        )
        if self.extra_pool_reduce == "mean":
            counts = h_other.new_zeros(batch_size, max_groups, 1)
            counts.scatter_add_(
                1,
                safe_group_id.unsqueeze(-1),
                sorted_mask.unsqueeze(-1).to(h_other.dtype),
            )
            pooled_h = pooled_h / counts.clamp_min(1.0)

        pooled_time.scatter_add_(
            1,
            safe_group_id,
            sorted_time * group_start.to(sorted_time.dtype),
        )
        group_count = group_start.sum(dim=1)
        arange_groups = torch.arange(max_groups, device=h_other.device)
        pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
        return pooled_h, pooled_time, pooled_mask

    def _forward_shared(
        self,
        event_seq: torch.LongTensor,
        time_seq: torch.FloatTensor,
        sex: torch.LongTensor,
        mode: str,
        padding_mask: torch.Tensor | None = None,
        t_query: torch.FloatTensor | None = None,
        other_type: torch.LongTensor | None = None,
        other_value: torch.Tensor | None = None,
        other_value_kind: torch.LongTensor | None = None,
        other_time: torch.FloatTensor | None = None,
        return_output: bool = False,
        **unused_kwargs,
    ) -> torch.Tensor | DeepHealthOutput:
        if unused_kwargs:
            unknown = ", ".join(sorted(unused_kwargs))
            raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
        if mode not in {"next_token", "all_future"}:
            raise ValueError("mode must be either 'next_token' or 'all_future'")
        if mode == "all_future" and t_query is None:
            raise ValueError("t_query is required when mode='all_future'")
        if (
            other_type is None
            or other_value is None
            or other_value_kind is None
            or other_time is None
        ):
            raise ValueError(
                "DeepHealth expects other_type, other_value, "
                "other_value_kind, and other_time."
            )

        if padding_mask is None:
            padding_mask = event_seq > PAD_IDX
        else:
            padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)

        event_len = event_seq.size(1)
        event_features = self.token_embedding(event_seq)
        event_time = time_seq

        if other_time.shape != other_type.shape:
            raise ValueError(
                "other_time must have the same shape as other_type, got "
                f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
            )
        other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
        other_features, other_mask = self.tokenizer(
            other_type=other_type,
            other_value=other_value,
            other_value_kind=other_value_kind,
        )
        other_features = other_features.to(device=event_seq.device)
        other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)

        event_features = torch.cat([event_features, other_features], dim=1)
        event_time = torch.cat([event_time, other_time], dim=1)
        event_valid_mask = torch.cat([padding_mask, other_mask], dim=1)

        batch_size = event_seq.size(0)
        sex_context = self.gender_embedding(sex)[:, None, :]
        event_features = (
            event_features
            + sex_context
            + self.age_encoding(event_time)
        )
        event_features = event_features * event_valid_mask.unsqueeze(-1).to(
            event_features.dtype
        )
        event_memory = self.event_norm(
            self.event_projection(event_features)
        )
        event_memory = event_memory * event_valid_mask.unsqueeze(-1).to(
            event_memory.dtype
        )

        if mode == "all_future":
            query_time = t_query[:, None]
            query_position = None
            query_features = (
                self.query_token.view(1, 1, -1)
                + sex_context
                + self.age_encoding(query_time)
            )
            query_valid_mask = torch.ones(
                batch_size,
                1,
                dtype=torch.bool,
                device=event_seq.device,
            )
        else:
            # Each event position is an independent parallel query. Including
            # its event feature preserves token-level next-step semantics.
            # Equal-time memory is additionally position-causal so a token
            # cannot read a later token that may be its Delphi2M target.
            query_time = event_time
            query_position = torch.arange(
                event_time.size(1),
                device=event_time.device,
            ).view(1, -1).expand(batch_size, -1)
            query_features = event_features
            query_valid_mask = event_valid_mask

        n_query = query_time.size(1)
        query_context = self.query_projection(query_features).reshape(
            batch_size,
            n_query,
            self.n_trajectory,
            self.trajectory_dim,
        )
        trajectory_state = (
            self.trajectory_prototypes.view(
                1,
                1,
                self.n_trajectory,
                self.trajectory_dim,
            )
            + query_context
        )
        event_invalid_mask = self._make_event_invalid_mask(
            event_valid_mask=event_valid_mask,
            event_time=event_time,
            query_time=query_time,
            query_position=query_position,
        )

        event_rope_cache = None
        query_rope_cache = None
        rbf_cache = None
        if self.time_mode == "relative":
            if self.rope is None or self.rbf is None:
                raise RuntimeError("Relative-time modules are not initialized")
            event_rope_cache = self.rope.precompute_cache(event_time)
            query_rope_cache = self.rope.precompute_cache(query_time)
            rbf_cache = self.rbf.precompute_cross_cache(
                query_time,
                event_time,
            )

        event_key_value = self.reasoning_core.project_event_memory(
            event_memory,
            event_rope_cache=event_rope_cache,
        )
        for _ in range(self.n_reasoning_rounds):
            trajectory_state = self.reasoning_core(
                trajectory_state=trajectory_state,
                event_key_value=event_key_value,
                event_invalid_mask=event_invalid_mask,
                query_rope_cache=query_rope_cache,
                rbf_cache=rbf_cache,
            )

        hidden_sequence = self.final_ln(
            trajectory_state.reshape(batch_size, n_query, self.d_model)
        )
        hidden_sequence = hidden_sequence * query_valid_mask.unsqueeze(-1).to(
            hidden_sequence.dtype
        )

        if mode == "all_future":
            hidden = hidden_sequence[:, 0, :]
            if return_output:
                return DeepHealthOutput(
                    hidden=hidden,
                    time_seq=t_query[:, None],
                    padding_mask=torch.ones(
                        hidden.size(0),
                        1,
                        dtype=torch.bool,
                        device=hidden.device,
                    ),
                    event_len=event_len,
                )
            return hidden
        if return_output:
            h_event = hidden_sequence[:, :event_len, :]
            t_event = event_time[:, :event_len]
            event_mask = event_valid_mask[:, :event_len]
            h_extra, t_extra, extra_mask = self._pool_other_by_time(
                h_other=hidden_sequence[:, event_len:, :],
                other_time=event_time[:, event_len:],
                other_mask=event_valid_mask[:, event_len:],
            )
            return DeepHealthOutput(
                hidden=torch.cat([h_event, h_extra], dim=1),
                time_seq=torch.cat([t_event, t_extra], dim=1),
                padding_mask=torch.cat([event_mask, extra_mask], dim=1),
                event_len=event_len,
            )
        return hidden_sequence[:, :event_len, :]

    def forward_next_token(self, **kwargs) -> torch.Tensor:
        return self._forward_shared(mode="next_token", **kwargs)

    def forward_all_future(self, **kwargs) -> torch.Tensor:
        return self._forward_shared(mode="all_future", **kwargs)

    def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor:
        mode = self.target_mode if target_mode is None else target_mode
        return self._forward_shared(mode=mode, **kwargs)

    def calc_risk(self, x: torch.Tensor) -> torch.Tensor:
        return self.risk_head(x)

    def calc_weibull_rho(self, x: torch.Tensor) -> torch.Tensor:
        if self.dist_mode != "weibull":
            raise RuntimeError(
                f"calc_weibull_rho called with dist_mode={self.dist_mode!r}"
            )
        return F.softplus(self.rho_head(x)) + 1e-6

    def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor:
        if self.dist_mode != "mixed":
            raise RuntimeError(
                f"calc_death_rho called with dist_mode={self.dist_mode!r}"
            )
        return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6