MLPNet

mlp_module.MLPNet · inherits nnx.Module

Full MLP: flatten [insample_y | futr_exog] -> ReLU'd Linear stack -> raw head. Mirrors the reference layer structure exactly: num_layers Linears (the first maps the flattened input to hidden_size, the rest are hidden-to-hidden), each followed by ReLU, then a separate un-activated out head. The flattened input is the scaled insample target of length input_size concatenated with the row-major flattened future-known exog window of shape [input_size + h, F] when futr_exog_size > 0.

__init__(self, *, h: int, input_size: int, futr_exog_size: int = 0, num_layers: int = 2, hidden_size: int = 1024, outputsize_multiplier: int = 1, rngs: nnx.Rngs)

(No prose summary provided in docstring.)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
futr_exog_size	int	0	(undocumented)
num_layers	int	2	(undocumented)
hidden_size	int	1024	(undocumented)
outputsize_multiplier	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

Raises: ValueError: if num_layers is less than 1.

__call__(self, insample_z: jnp.ndarray, futr_exog: jnp.ndarray | None = None) -> jnp.ndarray

(No prose summary provided in docstring.)

Parameter	Type	Default	Description
insample_z	jnp.ndarray	-	[B, L, 1] scaled target.
futr_exog	jnp.ndarray \| None	None	[B, L+h, F] or None.

Returns: jnp.ndarray (Shape [B, h, outputsize_multiplier]). (scaled space for point/quantile heads; raw pre-domain_map parameters for distribution heads).