StdMeanScaler

chronax.scaler.StdMeanScaler · inherits Scaler

Scales data to have zero mean and unit variance along a given dimension.

__init__(self, dim: int = -1, keepdim: bool = True, minimum_scale: float = 0.001)

(undocumented)

Parameter	Type	Default	Description
dim	int	-1	dimension along which to compute the scale
keepdim	bool	True	controls whether to retain dimension dim (of length 1) in the scale tensor, or suppress it.
minimum_scale	float	1e-3	default scale that is used for elements that are constantly zero along dimension dim.

__call__(self, data: torch.Tensor, padding_mask: torch.Tensor, weights: torch.Tensor, prefix_length: int | None = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Applies scaling to the input data.

Parameter	Type	Default	Description
data	torch.Tensor	-	(undocumented)
padding_mask	torch.Tensor	-	(undocumented)
weights	torch.Tensor	-	(undocumented)
prefix_length	int | None	None	(undocumented)

Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor] (The scaled data, location (mean), and scale (standard deviation)).

compute_causal_statistics

chronax.scaler.compute_causal_statistics

Compute causal mean and scale statistics along a specified dimension using a vectorized implementation of Welford's algorithm for numerical stability.

This implementation avoids explicit loops while maintaining the numerical stability of Welford's algorithm, achieving better performance with the same robustness against overflow issues.

Can optionally use global statistics to stabilize causal statistics by clamping extreme values, preventing instability while preserving a relaxed version of the causal property. This allows a controlled amount of future information leakage, introducing an explicit tradeoff between causality and stability. extreme values, preventing instability while preserving the causal property.

Parameter	Type	Default	Description
data	torch.Tensor	-	The input data tensor
weights	torch.Tensor	-	The weight tensor (same shape as data)
padding_mask	torch.Tensor	-	The padding mask tensor (same shape as data)
dim	int	-	The dimension along which to compute statistics (must be -1, the time dimension)
minimum_scale	float	-	Minimum scale value to use
use_bessel_correction	bool	True	Whether to use Bessel's correction to get an unbiased estimator
stabilize_with_global	bool	False	Whether to use global statistics to stabilize the causal statistics by clamping extreme values
scale_factor_exponent	float	10.0	Exponent that controls the allowed range of deviation from global scale. For example, with exponent=1.0, causal scale must be between 0.1x and 10x the global scale. With exponent=2.0, the range would be 0.01x to 100x.
prefix_length	int | None	None	If specified, the global statistics will be computed using only the prefix length requested. This is used for multistep decoding, where we only want to use the initial historical data to compute the global statistics. If stabilize_with_global is False, this parameter is ignored.

Returns: Tuple[torch.Tensor, torch.Tensor] (Causal mean and scale tensors, potentially stabilized with global statistics).

CausalStdMeanScaler

chronax.scaler.CausalStdMeanScaler · inherits Scaler

Causally scales the data along dimension dim which is expected to be the time dimension. For each position t along this dimension, the mean and standard deviation are computed using only data from positions up to t.

Can optionally stabilize causal statistics using global statistics to prevent extreme values, while preserving the causal property.

__init__(self, dim: int = -1, minimum_scale: float = 0.1, use_bessel_correction: bool = True, stabilize_with_global: bool = False, scale_factor_exponent: float = 10.0)

(undocumented)

Parameter	Type	Default	Description
dim	int	-1	dimension along which to compute the causal scale (must be -1, the last dimension)
minimum_scale	float	0.1	default scale that is used if the scale is below this threshold or for the first time step, since standard deviation cannot be computed with a single observation
use_bessel_correction	bool	True	whether to use Bessel's correction to get an unbiased estimator
stabilize_with_global	bool	False	whether to use global statistics to stabilize extreme causal statistics
scale_factor_exponent	float	10.0	exponent that controls the allowed range of deviation from global scale. For example, with exponent=1.0, causal scale must be between 0.1x and 10x the global scale. With exponent=2.0, the range would be 0.01x to 100x.

__call__(self, data: torch.Tensor, padding_mask: torch.Tensor, weights: torch.Tensor, prefix_length: int | None = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Applies causal scaling to the input data.

Parameter	Type	Default	Description
data	torch.Tensor	-	(undocumented)
padding_mask	torch.Tensor	-	(undocumented)
weights	torch.Tensor	-	(undocumented)
prefix_length	int | None	None	(undocumented)

Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor] (The scaled data, causal means, and causal scale).

CausalPatchStdMeanScaler

chronax.scaler.CausalPatchStdMeanScaler · inherits Scaler

Causally scales data in patches, where each patch uses statistics computed from all data up to and including that patch. Within each patch, all timesteps use the same scaling values.

This approach provides more stability than per-timestep causal scaling while still maintaining the causal property (not using future data).

The statistics are computed using Welford's algorithm, which provides better numerical stability compared to the direct computation of variance, especially when dealing with large values or a large number of data points.

__init__(self, dim: int = -1, patch_size: int = 32, minimum_scale: float = 0.1, use_bessel_correction: bool = True, stabilize_with_global: bool = False, scale_factor_exponent: float = 10.0)

(undocumented)

Parameter	Type	Default	Description
dim	int	-1	dimension along which to compute the causal scale. Must be -1 (the last dimension).
patch_size	int	32	number of timesteps in each patch
minimum_scale	float	0.1	default scale that is used for elements that are constantly zero along dimension dim or for the first patch.
use_bessel_correction	bool	True	whether to use Bessel's correction to get an unbiased estimator
stabilize_with_global	bool	False	whether to use global statistics to stabilize extreme causal statistics
scale_factor_exponent	float	10.0	exponent that controls the allowed range of deviation from global scale. For example, with exponent=1.0, causal scale must be between 0.1x and 10x the global scale. With exponent=2.0, the range would be 0.01x to 100x.

__call__(self, data: torch.Tensor, padding_mask: torch.Tensor, weights: torch.Tensor, prefix_length: int | None = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Applies causal patch scaling to the input data.

Parameter	Type	Default	Description
data	torch.Tensor	-	(undocumented)
padding_mask	torch.Tensor	-	(undocumented)
weights	torch.Tensor	-	(undocumented)
prefix_length	int | None	None	(undocumented)

Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor] (The scaled data, patch means, and patch scales).

scaler_types

chronax.scaler.scaler_types

Dictionary mapping string representations of scaler classes to their actual class definitions, used for deserialization of SafeTensors checkpoints.