chronax.models.timemixer.timemixer

`build_windows`

chronax.models.timemixer.timemixer_training.build_windows

Rolling windows over y, right-padded with h zeros.

`build_windows(y, input_size, h)`

Parameter	Type	Default	Description
`y`	`jnp.ndarray`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`h`	`int`	-	(undocumented)

Returns: tuple[jnp.ndarray, jnp.ndarray] (Returns (windows [n, input_size+h], target_mask [n, h]) with n = len(y) - input_size; the mask is 1.0 at real target positions and 0.0 in the zero-padded tail.)

`forward_loss`

chronax.models.timemixer.timemixer_training.forward_loss

Forward + masked point loss in ORIGINAL scale.

y_windows [B, input_size+h]; the univariate series carries a channel dim of 1 into the N-generic network.

`forward_loss(model, y_windows, target_mask=None, *, h, input_size, loss_fn)`

Parameter	Type	Default	Description
`model`	`TimeMixerNet`	-	(undocumented)
`y_windows`	`jnp.ndarray`	-	(undocumented)
`target_mask`	`jnp.ndarray \| None`	`None`	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`loss_fn`	`LossFn`	-	(undocumented)

Returns: jnp.ndarray

`train`

chronax.models.timemixer.timemixer_training.train

Train model in place via one nnx.scan. Returns per-step losses.

Window sampling replicates the reference's regime-dependent scheme: with fewer windows than windows_batch_size indices are drawn WITH replacement, otherwise a without-replacement permutation slice.

`train(model, y, *, h, input_size, max_steps, windows_batch_size, lr, seed, loss_fn)`

Parameter	Type	Default	Description
`model`	`TimeMixerNet`	-	(undocumented)
`y`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`max_steps`	`int`	-	(undocumented)
`windows_batch_size`	`int`	-	(undocumented)
`lr`	(undocumented)	-	(undocumented)
`seed`	`int`	-	(undocumented)
`loss_fn`	`LossFn`	-	(undocumented)

Returns: jnp.ndarray (per-step losses.) Raises: RuntimeError

`predict_step`

chronax.models.timemixer.timemixer_training.predict_step

Forecast next h steps from the final input_size of y. Returns [h, mult] in original scale (the net denormalizes internally).

`predict_step(model, y, *, h, input_size)`

Parameter	Type	Default	Description
`model`	`TimeMixerNet`	-	(undocumented)
`y`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)

Returns: jnp.ndarray ([h, mult] in original scale (the net denormalizes internally).)