Skip Connections
Skip Connections¤
Local operators used alongside the global spectral convolutions. They process high-frequency, localized information and act as residual connections to stabilize training.
Note: For the standard identity skip connection, neojax directly uses equinox.nn.Identity for simplicity.
neojax.nn.skip_connections.SoftGating
¤
Applies soft-gating by weighting the channels of the given input.
Given an input x of size (batch-size, channels, height, width),
this returns x * w
where w is of shape (1, channels, 1, 1)
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ndim
|
int
|
Number of spatial dimensions. |
required |
in_channels
|
int
|
Number of input channels. |
required |
out_channels
|
int | None
|
Number of output channels.
If provided, must match |
None
|
use_bias
|
bool
|
Whether to include a learnable bias. |
False
|
Internal Attributes
These fields store the internal layers state (and weights).
- weight (
Float[Array, ...]): Learnable channel-wise weights. - bias (
Float[Array, ...] | None): Optional learnable channel-wise bias.
__call__(x: Inexact[Array, 'c ...']) -> Inexact[Array, 'c ...']
¤
Applies soft-gating to input activations.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Inexact[Array, 'c ...']
|
Input activations. |
required |
Returns:
| Type | Description |
|---|---|
Inexact[Array, 'c ...']
|
Weighted activations. |
neojax.nn.skip_connections.Flattened1dConv
¤
Applies 1d convolution to flattened dimensions.
Skip layer that flattens all dimensions except for the leading channel dimension and applies 1d convolution over them, then un-flattens result.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
in_channels
|
int
|
Number of input channels. |
required |
out_channels
|
int
|
Number of output channels. |
required |
kernel_size
|
int
|
Size of the convolving kernel. |
required |
use_bias
|
bool
|
Whether to add a learnable bias to the output. |
False
|
Internal Attributes
These fields store the internal layers state (and weights).
- conv (
eqx.nn.Conv1d): The underlying 1D convolution layer. - out_channels (
int): Number of output channels.
__call__(x: Inexact[Array, 'in_c ...']) -> Inexact[Array, 'out_c ...']
¤
Applies 1d convolution to flattened dimensions.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Inexact[Array, 'in_c ...']
|
Input activations. |
required |
Returns:
| Type | Description |
|---|---|
Inexact[Array, 'out_c ...']
|
Output activations. |