The module loss.py provides masked point-loss functions for the Chronax TSMixer model. Pure JAX implementations of MAE and MSE with multiplicative masking, JIT- and grad-compatible. Shapes follow the [B, H, N] convention produced by TSMixer (batch, horizon, n_series), but the functions are generic and work for any matching broadcast-compatible shapes.

masked_mae

chronax.loss.masked_mae

Mean Absolute Error with multiplicative masking.

Returns: jnp.ndarray (The scalar loss value.)

masked_mse

chronax.loss.masked_mse

Mean Squared Error with multiplicative masking.

| Parameter | Type | Default | Description | |-----------|---------------|-------------| | y | jnp.ndarray | - | (undocumented) | | y_hat | jnp.ndarray | - | (undocumented) | | mask | Optional[jnp.ndarray] | None | (undocumented) | | horizon_weight | Optional[jnp.ndarray] | None | (undocumented) |

Returns: jnp.ndarray (The scalar loss value.)

resolve

chronax.loss.resolve

Return a callable loss from either a registry string or a callable.

Returns: Callable (The resolved loss function.)

Raises: * ValueError: If the provided string loss name is unknown.

LOSSES

chronax.loss.LOSSES

A dictionary mapping string aliases to loss functions.

Type: dict[str, Callable]

Contents:

{
    "mae": masked_mae,
    "mse": masked_mse,
}