Diagonal recurrent-state discretization and update

Your task

Function signature

diagonal_state_space_update(
    state: torch.Tensor,
    injection: torch.Tensor,
    log_delta: torch.Tensor,
    log_decay: torch.Tensor,
    input_matrix: torch.Tensor,
    norm_weight: torch.Tensor | None = None,
    norm_bias: torch.Tensor | None = None,
    update_mask: torch.Tensor | None = None,
    eps: float = 1e-05,
) -> dict[str, torch.Tensor]

Worked example

Decay and inject in different channels

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement one channel-diagonal state-space update. state and injection have the same floating shape [\ldots,C], with \mathrm{rank}\ge 1, C\ge 1 and all leading dimensions positive.
• log_delta [C], log_decay [C] and input_matrix [C,C] are floating, with the matrix indexed [\mathrm{output\_channel},\mathrm{input\_channel}]. Optional norm_weight [C] and norm_bias [C] default to ones and zeros. All floating values share \mathrm{float32}/64 dtype/device.
• eps is finite positive, default 10^{-5}, and is added INSIDE the square root of population variance for LayerNorm. Optional bool update_mask broadcasts to state.shape [:-1]; True updates a position, False returns its prior state unchanged.
• Compute \delta_{c}=\operatorname{softplus}(\mathrm{log\_delta\_c}) stably, \mathrm{decay\_c}=\operatorname{exp}(-\delta_{c}\cdot \operatorname{exp}(\mathrm{log\_decay\_c})), and \mathrm{discrete\_input\_matrix}[c,j]=\delta_{c}\cdot \mathrm{input\_matrix}[c,j].
• Normalize injection over its LAST C axis: \mathrm{centered}=\mathrm{injection}-\operatorname{mean}(\mathrm{injection}), \mathrm{normalized}=\mathrm{centered}/\sqrt{\operatorname{mean}(\mathrm{centered}^{2})+\varepsilon}, then multiply by norm_weight and add norm_bias when supplied.
• \mathrm{updated}[\ldots,c]=\mathrm{decay\_c}\cdot \mathrm{state}[\ldots,c]+\sum_{j} \mathrm{discrete\_input\_matrix}[c,j]\cdot \mathrm{normalized}[\ldots,j]. Apply update_mask only to the returned state.
• Return the updated-or-prior state, decay coefficients of shape (C,), and discrete input matrix of shape (C,C); discretized parameters are returned even if every position is masked. The diagonal refers to channelwise decay, not a diagonal of the residue-pair grid; the input matrix remains dense.
• log_delta is in [-40,40] and log_decay in [-20,20], and all required intermediate products/results are finite and representable. Very small decay may underflow to zero in the supplied dtype. Accept noncontiguous inputs, singleton channels and broadcast update masks; do not mutate inputs.
• This is one raw recurrence step, not a nonlinear folding layer or proof that the full recycled model converges.
• Compute softplus stably as \max(x,0)+\log(1+e^{-|x|}).

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def diagonal_state_space_update(
    state: torch.Tensor,  # (..., C)
    injection: torch.Tensor,  # (..., C)
    log_delta: torch.Tensor,  # (C,)
    log_decay: torch.Tensor,  # (C,)
    input_matrix: torch.Tensor,  # (C, C)
    norm_weight: torch.Tensor | None = None,  # (C,)
    norm_bias: torch.Tensor | None = None,  # (C,)
    update_mask: torch.Tensor | None = None,  # Any shape broadcastable to state.shape(:-1).
    eps: float = 1e-05,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    state: (..., C).
    decay: (C,).
    discrete_input_matrix: (C, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement diagonal_state_space_update")

Sources and further reading

• www.biorxiv.org
• huggingface/transformers · modeling_esmfold2.py
• huggingface/transformers · modeling_esmfold2.py