eTimeCNNEncoder#

class eTimeCNNEncoder(in_rep, out_rep, hidden_channels, time_horizon, activation=ReLU(), batch_norm=False, bias=True, mlp_hidden=(128,), downsample='stride', append_last_frame=False, init_scheme='xavier_uniform')[source]#

Bases: Module

Equivariant 1D CNN encoder built from channel-equivariant blocks.

Inputs are plain tensors of shape (N, in_rep.size, H). Each conv block halves the time horizon via stride-2 convolution; optional eRMSNorm and pointwise activation follow. The flattened feature map feeds either an equivariant head (eMLP) or an invariant head (iMLP) depending on whether out_rep contains only the trivial irrep.

The encoder defines:

\[\mathbf{f}_{\mathbf{\theta}}: \mathcal{X}^{H} \to \mathcal{Y},\]

where \(H\) is the input horizon, \(\mathcal{X}\) is the channel feature space transforming by \(\rho_{\mathcal{X}}\), and \(\mathcal{Y}\) transforms by \(\rho_{\mathcal{Y}}\).

Functional constraint (equivariant head):

\[\mathbf{f}_{\mathbf{\theta}}(\rho_{\mathcal{X}}(g)\mathbf{x}) = \rho_{\mathcal{Y}}(g)\mathbf{f}_{\mathbf{\theta}}(\mathbf{x}) \quad \forall g\in\mathbb{G}.\]

If out_rep is trivial-only, the head is invariant:

\[\mathbf{f}_{\mathbf{\theta}}(\rho_{\mathcal{X}}(g)\mathbf{x}) = \mathbf{f}_{\mathbf{\theta}}(\mathbf{x}) \quad \forall g\in\mathbb{G}.\]

Create an equivariant time-series CNN encoder.

Parameters:

  • in_rep (Representation) – Input representation \(\rho_{\text{in}}\) defining the group action on the input channels.

  • out_rep (Representation) – Output representation \(\rho_{\text{out}}\). If it contains only trivial irreps, an iMLP head is used; otherwise an eMLP head is used.

  • hidden_channels (list[int]) – List of output channel counts for each convolution block.

  • time_horizon (int) – Length of the input time series (number of frames).

  • activation (Module) – Non-linearity applied after every convolution block.

  • batch_norm (bool) – Whether to include channel-wise RMS normalization.

  • bias (bool) – Whether to include bias in convolutions and linear heads.

  • mlp_hidden (list[int]) – Hidden layer widths for the final MLP head.

  • downsample (str) – Downsampling strategy, either 'stride' (stride-2 conv) or 'pooling' (max pool).

  • append_last_frame (bool) – Whether to concatenate the last frame of the input to the encoding before the head.

  • init_scheme (str | None) – Initialization scheme for equivariant layers.

check_equivariance(atol=1e-05, rtol=1e-05)[source]#

Check equivariance under channel actions of the underlying group.

Parameters:
forward(x)[source]#

Encode input of shape (N, in_rep.size, H) into (N, out_rep.size).

Return type:

Tensor

Parameters:

x (Tensor)