CausalConv1d

tcn_module.CausalConv1d

Causal dilated 1-D convolution followed by an activation. Operates in channel-first layout [B, C, L]; causality is realized as K left-shifted matmul taps (see __call__). Weight shape is (out_channels, in_channels, kernel_size).

__init__(self, in_channels: int, out_channels: int, kernel_size: int, padding: int, dilation: int, activation: str, *, rngs: nnx.Rngs)

(undocumented)

Parameter	Type	Default	Description
in_channels	int	-	(undocumented)
out_channels	int	-	(undocumented)
kernel_size	int	-	(undocumented)
padding	int	-	(undocumented)
dilation	int	-	(undocumented)
activation	str	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray) -> jnp.ndarray

x: [B, C_in, L] -> [B, C_out, L].

Computed as K shifted matmuls (out[t] = Σ_k W[:,:,k] · x[t-(K-1-k)·d]) rather than lax.conv_general_dilated: the XLA CPU backend pessimizes conv primitives inside the lax.scan training loop, while matmuls keep their fast dot path there. K and the shift amounts are static ctor config, so this traces unchanged under jit/vmap.

Parameters:

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

Returns: jnp.ndarray

TemporalConvolutionEncoder

tcn_module.TemporalConvolutionEncoder

Sequential stack of causal dilated convs. Layer i uses padding = (kernel_size-1) * dilations[i]; layer 0 maps in_channels -> out_channels, the rest out_channels -> out_channels. Input/output are time-first [B, L, C] (transposed to channel-first internally).

__init__(self, in_channels: int, out_channels: int, kernel_size: int, dilations: tuple, activation: str = 'ReLU', *, rngs: nnx.Rngs)

(undocumented)

Parameter	Type	Default	Description
in_channels	int	-	(undocumented)
out_channels	int	-	(undocumented)
kernel_size	int	-	(undocumented)
dilations	tuple	-	(undocumented)
activation	str	"ReLU"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray) -> jnp.ndarray

x: [B, L, C_in] -> [B, L, C_out].

Parameters:

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

Returns: jnp.ndarray

MLP

tcn_module.MLP

MLP decoder head: num_layers Linears total, ReLU between. num_layers=1 is a direct linear projection. For num_layers>=2: input Linear, num_layers-2 hidden Linears, output Linear — ReLU after every layer but the last. Dropout is omitted (TCN runs it at 0.0).

__init__(self, in_features: int, out_features: int, hidden_size: int, num_layers: int, *, rngs: nnx.Rngs)

(undocumented)

Parameter	Type	Default	Description
in_features	int	-	(undocumented)
out_features	int	-	(undocumented)
hidden_size	int	-	(undocumented)
num_layers	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray) -> jnp.ndarray

(undocumented)

Parameters:

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

Returns: jnp.ndarray

TCNNet

tcn_module.TCNNet

Full TCN: encoder -> context adapter -> futr residual concat -> MLP decoder.

Forward: 1. concat scaled insample_y with the historic slice of future-known exog -> encoder input [B, L, 1+F]; 2. TemporalConvolutionEncoder -> [B, L, C]; 3. transpose to [B, C, L], context_adapter = Linear(L -> h) over the time axis -> [B, C, h] (natively handles h > input_size); 4. concat the horizon slice of futr exog as extra channels -> [B, C+F, h]; 5. transpose to [B, h, C+F], per-timestep MLP decoder -> [B, h, mult].

The forward is fully deterministic (no dropout or batchnorm), so no RNG or train/eval mode flags are needed.

__init__(self, *, h: int, input_size: int, kernel_size: int = 2, dilations: tuple = (1, 2, 4, 8, 16), encoder_hidden_size: int = 128, encoder_activation: str = 'ReLU', decoder_hidden_size: int = 128, decoder_layers: int = 2, futr_exog_size: int = 0, outputsize_multiplier: int = 1, rngs: nnx.Rngs)

(undocumented)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
kernel_size	int	2	(undocumented)
dilations	tuple	(1, 2, 4, 8, 16)	(undocumented)
encoder_hidden_size	int	128	(undocumented)
encoder_activation	str	"ReLU"	(undocumented)
decoder_hidden_size	int	128	(undocumented)
decoder_layers	int	2	(undocumented)
futr_exog_size	int	0	(undocumented)
outputsize_multiplier	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

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

insample_z: [B, L, 1] scaled target; futr_exog: [B, L+h, F] or None.

Returns [B, h, outputsize_multiplier] in scaled space.

Parameters:

Parameter	Type	Default	Description
insample_z	jnp.ndarray	-	(undocumented)
futr_exog	jnp.ndarray | None	None	(undocumented)

Returns: jnp.ndarray (Returns [B, h, outputsize_multiplier] in scaled space.)