CustomConv1d

bitcn_module.CustomConv1d · inherits nnx.Module

Forward- or backward-looking causal dilated 1-D convolution. Mirrors NF CustomConv1d. Operates in torch's channel-first layout [B, C, L] via jax.lax.conv_general_dilated with dimension numbers ('NCH', 'OIH', 'NCH') and a (low, high) causal pad: mode="backward" pads on the left (looks into the past), mode="forward" pads on the right. groups=1 throughout (as in NF). Weight shape is torch's (out_channels, in_channels, kernel_size).

__init__(self, in_channels: int, out_channels: int, kernel_size: int, padding: int = 0, dilation: int = 1, mode: str = 'backward', *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
in_channels	int	-	(undocumented)
out_channels	int	-	(undocumented)
kernel_size	int	-	(undocumented)
padding	int	0	(undocumented)
dilation	int	1	(undocumented)
mode	str	"backward"	Must be 'backward' or 'forward'.
rngs	nnx.Rngs	-	(undocumented)

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

Performs the convolution.

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input tensor: [B, C_in, L].

Returns: jnp.ndarray (Output tensor: [B, C_out, L]).

TCNCell

bitcn_module.TCNCell · inherits nnx.Module

One Temporal Convolutional Network cell (NF TCNCell). conv1 is a dilated causal conv (kernel_size=2); conv2 is a pointwise (kernel_size=1) conv widening to 2 * in_channels whose output is split into a residual update h_next and a skip contribution out_next.

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

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

__call__(self, h_prev: jnp.ndarray, out_prev: jnp.ndarray, deterministic: bool)

Performs the TCN cell operation, returning updated hidden state and skip accumulator.

Parameter	Type	Default	Description
h_prev	jnp.ndarray	-	(undocumented)
out_prev	jnp.ndarray	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: tuple[jnp.ndarray, jnp.ndarray] (Updated hidden state, updated skip accumulator).

BiTCNNet

bitcn_module.BiTCNNet · inherits nnx.Module

Univariate BiTCN backbone: [B, input_size, 1] -> [B, h, 1]. No-exogenous case (backward TCN only). Identity scaling (NF default): the network operates in raw scale. The number of backward layers is set so the dilated receptive field covers the lookback window. During training pass deterministic=False; the nnx.Rngs dropout stream supplies noise.

__init__(self, *, h: int, input_size: int, hidden_size: int, dropout: float, rngs: nnx.Rngs)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
hidden_size	int	-	(undocumented)
dropout	float	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

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

Performs the forward pass through the BiTCN network.

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input sequence [B, L, 1].
deterministic	bool	-	If True, disables dropout.

Returns: jnp.ndarray (Forecast output [B, h, 1]).