RevIN#

class sdm.nn.RevIN(affine_channels: int | None = None, eps: float = 1e-05, device: device | str | None = None, dtype: dtype | None = None)#

Bases: Module

Reversible instance normalization from the “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift” paper.

Parameters:

  • affine_channels (int | None) – The number of channels for learned affine parameters. If None, no affine transformation is applied.

  • eps (float) – The value added to the supplied scale for numerical stability.

  • device (device | str | None) – The device.

  • dtype (dtype | None) – The dtype.

static stats(x: Tensor, mask: Tensor | None = None, correction: int = 1) → tuple[Tensor, Tensor]#

Compute mean and standard deviation over the final dimension.

NaNs and masked values are excluded. Mean and standard deviation are computed and returned in torch.float64.

Parameters:

  • x (Tensor) – Tensor with shape [..., R, D].

  • mask (Tensor | None) – Boolean mask with shape [..., R, D]. Values marked True are excluded.

  • correction (int) – Difference between the sample size and sample degrees of freedom.

Returns:

Mean and standard deviation with shape [..., R, 1].

Return type:

tuple[Tensor, Tensor]

static running_stats(x: Tensor, mask: Tensor | None = None, state: tuple[Tensor, Tensor, Tensor] | None = None, correction: int = 1) → tuple[Tensor, Tensor, Tensor]#

Compute cumulative statistics over patched sequences.

NaNs and masked values are excluded. Mean and standard deviation are computed and returned in torch.float64.

Parameters:

  • x (Tensor) – Tensor with shape [..., R, D].

  • mask (Tensor | None) – Boolean mask with shape [..., R, D]. Values marked True are excluded.

  • state (tuple[Tensor, Tensor, Tensor] | None) – Previous mean, standard deviation, and count for streaming statistics.

  • correction (int) – Degrees of freedom subtracted from the observed count.

Returns:

Cumulative mean, standard deviation, and observed count for the supplied sequence, each with shape [..., R, 1].

Return type:

tuple[Tensor, Tensor, Tensor]

forward(x: Tensor, mean: Tensor, scale: Tensor) → Tensor#

The forward pass.

Parameters:

  • x (Tensor) – Tensor with shape [..., R, D].

  • mean (Tensor) – The mean with shape [..., R, 1].

  • scale (Tensor) – The scale with shape [..., R, 1].

Return type:

Tensor

inverse(x: Tensor, mean: Tensor, scale: Tensor) → Tensor#

Invert the forward pass.

Parameters:

  • x (Tensor) – Tensor with shape [..., R, D].

  • mean (Tensor) – The mean with shape [..., R, 1].

  • scale (Tensor) – The scale with shape [..., R, 1].

Return type:

Tensor