mae

chronax.mlp_losses.mae

Mean absolute error.

Parameter	Type	Default	Description
pred	jnp.ndarray	-	(undocumented)
target	jnp.ndarray	-	(undocumented)

mse

chronax.mlp_losses.mse

Mean squared error.

Parameter	Type	Default	Description
pred	jnp.ndarray	-	(undocumented)
target	jnp.ndarray	-	(undocumented)

huber

chronax.mlp_losses.huber

Huber loss, delta = 1.0.

Parameter	Type	Default	Description
pred	jnp.ndarray	-	(undocumented)
target	jnp.ndarray	-	(undocumented)

weighted_average

chronax.mlp_losses.weighted_average

Weighted average with zero-weight masking, matching neuralforecast's weighted_average: cells with weight 0 contribute exactly 0 (so garbage or NaN values under a zero mask never leak in), and the denominator is max(sum(weights), 1.0) — clamped to one, not to an epsilon.

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
weights	jnp.ndarray \| None	None	(undocumented)
axis	-	None	(undocumented)

outputsize_multiplier

chronax.mlp_losses.outputsize_multiplier

Output-head width implied by loss (1 for point losses).

Parameter	Type	Default	Description
loss	-	-	(undocumented)

resolve

chronax.mlp_losses.resolve

Return a loss object from a registry string, a callable, an MQ instance, or a GMM instance.

Parameter	Type	Default	Description
loss	str \| LossFn \| MultiQuantileLoss \| GMM	-	(undocumented)

MultiQuantileLoss

chronax.mlp_losses.MultiQuantileLoss

Multi-quantile (pinball) loss. __call__(pred[...,h,Q], target[...,h]).

QL(y, y_hat, q) = q*(y-y_hat)+ + (1-q)*(y_hat-y)+, averaged over quantiles and all elements. Quantiles are sorted, must lie in (0, 1), and must include 0.5 (the median / "mean" head). Picklable (holds a plain tuple).

__init__(self, quantiles=(0.1, 0.5, 0.9))

Parameter	Type	Default	Description
quantiles	Sequence[float]	(0.1, 0.5, 0.9)	(undocumented)

__call__(self, pred, target) -> jnp.ndarray

Parameter	Type	Default	Description
pred	jnp.ndarray	-	(undocumented)
target	jnp.ndarray	-	(undocumented)

GMM

chronax.mlp_losses.GMM

Gaussian Mixture Model distribution loss (port of neuralforecast GMM).

The network head emits (2 + weighted) * n_components values per horizon step: component means and (pre-softplus) standard deviations, plus mixture weight logits when weighted=True; otherwise weights are uniform 1/K. __call__ is the mixture negative log-likelihood -logsumexp(log N(y; mu_k, sigma_k) + log w_k) reduced by :func:weighted_average over the mask. batch_correlation / horizon_correlation sum the component log-likelihoods over the batch / horizon axis before mixing (composite-likelihood variants; they assume [batch, horizon, components] inputs).

scale_decouple maps scaled-space parameters back to the data scale (means*scale + loc, stds = (softplus(stds) + eps) * scale): the network optimizes in the scaler's normalized space while the likelihood is evaluated in original units. The eps floor applies whenever an anchor is given — the identity scaler anchors with real 0/1 tensors, so the floor is active there too, matching the reference exactly.

sample draws num_samples Monte-Carlo paths directly from the mixture with an explicit PRNG key. The torch reference reaches the same distribution through an evenly-spaced quantile grid that is then bootstrap-resampled with an unseeded generator (a workaround for its column-oriented predict path); direct seeded draws are equal in distribution and reproducible.

The point forecast surface uses :meth:analytic_mean (the exact mixture mean) rather than the Monte-Carlo mean of the samples.

return_params is accepted for API parity but not supported.

__init__(self, n_components=1, level=(80, 90), quantiles=None, num_samples=1000, return_params=False, batch_correlation=False, horizon_correlation=False, weighted=False)

Parameter	Type	Default	Description
n_components	int	1	(undocumented)
level	Sequence[float]	(80, 90)	(undocumented)
quantiles	Sequence[float] \| None	None	(undocumented)
num_samples	int	1000	(undocumented)
return_params	bool	False	(undocumented)
batch_correlation	bool	False	(undocumented)
horizon_correlation	bool	False	(undocumented)
weighted	bool	False	(undocumented)

domain_map(self, output) -> tuple

Split the raw head [..., (2+weighted)*K] into equal parameter chunks (means, stds[, weight_logits]) of [..., K] each.

Parameter	Type	Default	Description
output	jnp.ndarray	-	(undocumented)

scale_decouple(self, output, loc=None, scale=None, eps=0.2) -> tuple

Positivity-map the stds (softplus) and, when an anchor is given, map parameters to the data scale. loc/scale must broadcast against the [..., h, K] parameter arrays (e.g. [B, 1, 1]).

Parameter	Type	Default	Description
output	tuple	-	(undocumented)
loc	-	None	(undocumented)
scale	-	None	(undocumented)
eps	float	0.2	(undocumented)

__call__(self, y, distr_args, mask=None) -> jnp.ndarray

Masked negative log-likelihood. y [..., h] in the same space as the (decoupled) distr_args; mask weights positions per :func:weighted_average.

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
distr_args	tuple	-	(undocumented)
mask	jnp.ndarray \| None	None	(undocumented)

analytic_mean(self, distr_args) -> jnp.ndarray

Exact mixture mean sum_k w_k mu_k over the component axis.

Parameter	Type	Default	Description
distr_args	tuple	-	(undocumented)

sample(self, distr_args, num_samples=None, *, key) -> jnp.ndarray

Draw [..., num_samples] Monte-Carlo paths from the mixture.

Parameter	Type	Default	Description
distr_args	tuple	-	(undocumented)
num_samples	int \| None	None	(undocumented)
key	-	-	(undocumented)